#ifndef UNICODE
#define UNICODE
#endif
#ifndef UNICODE
#define _UNICODE
#endif

#include <stdio.h>
#include <windows.h>
#include <wchar.h>
#include <string.h>
#include <ntsecapi.h> // for LsaNtStatusToWinError()
#pragma hdrstop



#define IN
#define OUT


//
// Define the file information class values
//
// WARNING:  The order of the following values are assumed by the I/O system.
//           Any changes made here should be reflected there as well.
//

typedef enum _FILE_INFORMATION_CLASS {
	FileDirectoryInformation = 1,
	FileFullDirectoryInformation,
	FileBothDirectoryInformation,
	FileBasicInformation,
	FileStandardInformation,
	FileInternalInformation,
	FileEaInformation,
	FileAccessInformation,
	FileNameInformation,
	FileRenameInformation,
	FileLinkInformation,
	FileNamesInformation,
	FileDispositionInformation,
	FilePositionInformation,
	FileFullEaInformation,
	FileModeInformation,
	FileAlignmentInformation,
	FileAllInformation,
	FileAllocationInformation,
	FileEndOfFileInformation,
	FileAlternateNameInformation,
	FileStreamInformation,
	FilePipeInformation,
	FilePipeLocalInformation,
	FilePipeRemoteInformation,
	FileMailslotQueryInformation,
	FileMailslotSetInformation,
	FileCompressionInformation,
	FileCopyOnWriteInformation,
	FileCompletionInformation,
	FileMoveClusterInformation,
	FileOleClassIdInformation,
	FileOleStateBitsInformation,
	FileNetworkOpenInformation,
	FileObjectIdInformation,
	FileOleAllInformation,
	FileOleDirectoryInformation,
	FileContentIndexInformation,
	FileInheritContentIndexInformation,
	FileOleInformation,
	FileMaximumInformation
} FILE_INFORMATION_CLASS, *PFILE_INFORMATION_CLASS;

//
// Define the various structures which are returned on query operations
//

typedef struct _FILE_BASIC_INFORMATION {
	LARGE_INTEGER CreationTime;
	LARGE_INTEGER LastAccessTime;
	LARGE_INTEGER LastWriteTime;
	LARGE_INTEGER ChangeTime;
	ULONG FileAttributes;
} FILE_BASIC_INFORMATION, *PFILE_BASIC_INFORMATION;

typedef struct _FILE_STANDARD_INFORMATION {
	LARGE_INTEGER AllocationSize;
	LARGE_INTEGER EndOfFile;
	ULONG NumberOfLinks;
	BOOLEAN DeletePending;
	BOOLEAN Directory;
} FILE_STANDARD_INFORMATION, *PFILE_STANDARD_INFORMATION;

typedef struct _FILE_POSITION_INFORMATION {
	LARGE_INTEGER CurrentByteOffset;
} FILE_POSITION_INFORMATION, *PFILE_POSITION_INFORMATION;

typedef struct _FILE_ALIGNMENT_INFORMATION {
	ULONG AlignmentRequirement;
} FILE_ALIGNMENT_INFORMATION, *PFILE_ALIGNMENT_INFORMATION;

typedef struct _FILE_NETWORK_OPEN_INFORMATION {
	LARGE_INTEGER CreationTime;
	LARGE_INTEGER LastAccessTime;
	LARGE_INTEGER LastWriteTime;
	LARGE_INTEGER ChangeTime;
	LARGE_INTEGER AllocationSize;
	LARGE_INTEGER EndOfFile;
	ULONG FileAttributes;
} FILE_NETWORK_OPEN_INFORMATION, *PFILE_NETWORK_OPEN_INFORMATION;

typedef struct _FILE_DISPOSITION_INFORMATION {
	BOOLEAN DeleteFile;
} FILE_DISPOSITION_INFORMATION, *PFILE_DISPOSITION_INFORMATION;

typedef struct _FILE_END_OF_FILE_INFORMATION {
	LARGE_INTEGER EndOfFile;
} FILE_END_OF_FILE_INFORMATION, *PFILE_END_OF_FILE_INFORMATION;


typedef struct _FILE_FULL_EA_INFORMATION {
	ULONG NextEntryOffset;
	UCHAR Flags;
	UCHAR EaNameLength;
	USHORT EaValueLength;
	CHAR EaName[1];
} FILE_FULL_EA_INFORMATION, *PFILE_FULL_EA_INFORMATION;


struct FILE_STREAM_INFORMATION
{
	ULONG NextEntryOffset;
	ULONG StreamNameLength;
	LARGE_INTEGER StreamSize;
	LARGE_INTEGER StreamAllocationSize;
	WCHAR StreamName[1];
};


//
// Define the base asynchronous I/O argument types
//

typedef struct _IO_STATUS_BLOCK {
	NTSTATUS Status;
	ULONG Information;
} IO_STATUS_BLOCK, *PIO_STATUS_BLOCK;


DWORD __stdcall NtQueryInformationFile(
	IN HANDLE FileHandle,
	OUT PIO_STATUS_BLOCK IoStatusBlock,
	OUT PVOID FileInformation,
	IN ULONG Length,
	IN FILE_INFORMATION_CLASS FileInformationClass
	);

typedef DWORD (__stdcall *NQIF)(
	IN HANDLE FileHandle,
	OUT PIO_STATUS_BLOCK IoStatusBlock,
	OUT PVOID FileInformation,
	IN ULONG Length,
	IN FILE_INFORMATION_CLASS FileInformationClass
	);



NQIF nqif = NULL;
const int MAXLEN = 2048; // arbitrary filename length limit



void checkOneFile( const wchar_t *file, int& numFilesWithStreams, int& numStreams )
{
	HANDLE h;
	DWORD rc;
	IO_STATUS_BLOCK iosb;
	static byte fsibuf[32768];
	FILE_STREAM_INFORMATION *fsi;
	bool mustPrintName = true;
	static const wchar_t * const DefaultStreamName = L"::$DATA";
	static const ULONG DefaultStreamNameLength = 7;

	// note that while we do have to open the file, not even read access is needed
	h = CreateFile( file, 0, FILE_SHARE_READ | FILE_SHARE_WRITE, NULL,
		OPEN_EXISTING, 0, NULL );
	if ( h == INVALID_HANDLE_VALUE )
	{
		wprintf( L"Error %d opening \"%s\".\n", GetLastError(), file );
		return;
	}

	rc = nqif( h, &iosb, &fsibuf, sizeof( fsibuf ), FileStreamInformation );
	if ( rc != 0 )
	{
		wprintf( L"Error %d [NTSTATUS %ld/%08lxh] on NQIF() of \"%s\".\n",
			LsaNtStatusToWinError( rc ), rc, rc, file );
		CloseHandle( h );
		return;
	}
	
	CloseHandle( h );

	fsi = (FILE_STREAM_INFORMATION *) fsibuf;

	while ( 1 )
	{
		if ( fsi->StreamNameLength != 2 * DefaultStreamNameLength ||
			wcsncmp( fsi->StreamName, DefaultStreamName, DefaultStreamNameLength ) != 0 )
		{
			// non-default stream found!
			if ( mustPrintName )
			{
				// first stream for this file, bump file count and show file name
				++ numFilesWithStreams;
				wprintf( L"\n%s\n", file );
				mustPrintName = false;
			}
			// bump stream count, show stream info
			++ numStreams;
			wprintf( L"  %8I64u  %-*.*s\n", fsi->StreamSize, fsi->StreamNameLength / 2,
				fsi->StreamNameLength / 2, fsi->StreamName );
		}
		if ( fsi->NextEntryOffset == 0 )
			break;
		fsi = (FILE_STREAM_INFORMATION *) ( fsi->NextEntryOffset + (byte *) fsi );
	}

	return;
}



void doFilespec( const wchar_t *srchPath, int& numFilesWithStreams, int& numStreams )
{
	wchar_t path[MAXLEN] = L"";
	wchar_t *p, *endOfPath;
	WIN32_FIND_DATA wfd;
	HANDLE ffh;

	if ( srchPath == NULL || *srchPath == L'\0' || wcslen( srchPath ) >= MAXLEN )
		return;

	wcscpy( path, srchPath );
	for ( p = path; *p != L'\0'; ++ p )
	{
		if ( *p == L'/' )
			*p = L'\\';
	}

	p = wcsrchr( path, L'\\' ); // isolate last backslash
	if ( p == NULL ) // no backslash, only a filespec
		endOfPath = path;
	else
	{
		endOfPath = p + 1;
		if ( *endOfPath == L'\0' ) // no filespec?
			wcscpy( endOfPath, L"*" ); // invent one
	}

	// loop over files
	ffh = FindFirstFile( path, &wfd );
	if ( ffh != INVALID_HANDLE_VALUE )
	{
		BOOL ok;
		do
		{
			if ( ( wfd.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY ) != 0 )
			{
				// it's a directory -- may want to recurse
			}
			else // it's a file
			{
				wcscpy( endOfPath, wfd.cFileName ); // tack file onto path
				checkOneFile( path, numFilesWithStreams, numStreams);
			}
			ok = FindNextFile( ffh, &wfd );
		} while ( ok );
		FindClose( ffh );
	}
	else
		wprintf( L"Error %d when searching for \"%s\".\n", GetLastError(), path );

	return;
}



extern "C"
int __cdecl wmain( int argc, wchar_t *argv[], wchar_t *envp[] )
{
	int numFilesWithStreams = 0, numStreams = 0, i;
	HINSTANCE hNtdll;

	if ( argc == 1 )
	{
		_putws( L"usage: streams filespec [...]" );
		_putws( L"To specify all files in a directory, use \"dir\\\" or \"dir\\*\"." );
		_putws( L"streams _only_ reports on files having non-default streams" );
		return 1;
	}

	hNtdll = LoadLibrary( L"ntdll.dll" );
	if ( hNtdll < (HINSTANCE) 33 )
	{
		wprintf( L"Error %d [%d] loading NTDLL.DLL. Huh?\n", (int) hNtdll, GetLastError() );
		return -2;
	}

	nqif = (NQIF) GetProcAddress( hNtdll, "NtQueryInformationFile" );
	if ( nqif == NULL )
	{
		wprintf( L"NQIF() missing from NTDLL.DLL. Huh?\n" );
		FreeLibrary( hNtdll );
		return -2;
	}

	for ( i = 1; i < argc; ++ i )
		doFilespec( argv[i], numFilesWithStreams, numStreams );

	FreeLibrary( hNtdll );

	return numStreams;
}