/******************************************************************************
fgdump - by fizzgig and the foofus.net group
Copyright (C) 2005 by fizzgig
http://www.foofus.net

This program is free software; you can redistribute it and/or
modify it under the terms of the GNU General Public License
as published by the Free Software Foundation; either version 2
of the License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
******************************************************************************/

#include "StdAfx.h"
#include "sharefinder.h"
#include <Lm.h>

ShareFinder::ShareFinder()
{

}

ShareFinder::~ShareFinder()
{

}

bool ShareFinder::EnumerateShares(char* szServer)
{
	PSHARE_INFO_502 BufPtr, p;
	NET_API_STATUS res;
	DWORD er=0, tr=0, resume=0, i;
	wchar_t server[MAX_PATH];
	char szServerWithSlashes[MAX_PATH];

	::ZeroMemory(server, MAX_PATH);
	::ZeroMemory(szServerWithSlashes, MAX_PATH);
	_snprintf(szServerWithSlashes, MAX_PATH, "\\\\%s", szServer);
	mbstowcs(server, szServerWithSlashes, strlen(szServerWithSlashes));

	do
	{
		// Fuck Microsoft and it's lame-ass unicode crap
		res = NetShareEnum((LPSTR)server, 502, (LPBYTE*)&BufPtr, -1, &er, &tr, &resume);
		if(res == ERROR_SUCCESS || res == ERROR_MORE_DATA)
		{
			p = BufPtr;
			for(i=1; i <= er; i++)
			{
				Log.ReportError(INFO, "Found share %S, whose physical path is %S\n", p->shi502_netname, p->shi502_path);
				p++;
			}

			NetApiBufferFree(BufPtr);
		}
		else 
			Log.ReportError(CRITICAL, "EnumerateShares returned an error of %ld\n",res);
	}
	while (res == ERROR_MORE_DATA); // end do

	return true;
}

bool ShareFinder::BindUploadShareToLocalDrive(char* szServer, char* lpszLocalDrive, int nBufferSize, char** lplpPhysicalPath)
{
	PSHARE_INFO_502 BufPtr, p;
	NET_API_STATUS res;
	DWORD er=0, tr=0, resume=0, i;
	wchar_t server[MAX_PATH];
	char szTemp[MAX_PATH];
	bool bBound = false;
	char szServerWithSlashes[MAX_PATH];

	::ZeroMemory(server, MAX_PATH);
	::ZeroMemory(szServerWithSlashes, MAX_PATH);
	::ZeroMemory(*lplpPhysicalPath, nBufferSize);
	_snprintf(szServerWithSlashes, MAX_PATH, "\\\\%s", szServer);
	mbstowcs(server, szServerWithSlashes, strlen(szServerWithSlashes));

	do
	{
		// Fuck Microsoft and it's lame-ass unicode crap
		res = NetShareEnum((LPSTR)server, 502, (LPBYTE*)&BufPtr, -1, &er, &tr, &resume);
		if(res == ERROR_SUCCESS || res == ERROR_MORE_DATA)
		{
			p = BufPtr;
			for(i=1; i <= er; i++)
			{
				::ZeroMemory(szTemp, MAX_PATH);
				wcstombs(szTemp, (LPWSTR)(p->shi502_netname), MAX_PATH);

				// Look for shares that are not SYSVOL or NETLOGON, and that have a physical path
				if (stricmp(szTemp, "SYSVOL") != 0 && stricmp(szTemp, "NETLOGON") != 0 && wcslen((LPWSTR)(p->shi502_path)) > 0)
				{
					// If this is a potentially workable share, bind the drive and try uploading something
					if (BindDrive(lpszLocalDrive, szServerWithSlashes, szTemp))
					{
						// Success!
						// Copy the physical path to the out variable
						wcstombs(szTemp, (LPWSTR)(p->shi502_path), MAX_PATH);
						strncpy(*lplpPhysicalPath, szTemp, nBufferSize);
						bBound = true;
						break;
					}
					// Otherwise continue and try another share
				}
				
				p++;
			}

			NetApiBufferFree(BufPtr);
		}
		else 
			Log.ReportError(CRITICAL, "BindUploadShareToLocalDrive returned an error of %ld\n",res);
	}
	while (res == ERROR_MORE_DATA); // end do

	return true;
}

bool ShareFinder::BindDrive(char* szDrive, char* szServer, char* szShare)
{
	DWORD dwResult; 
	NETRESOURCE nr; 
	char szTemp[MAX_PATH];
	char szTempFilename[MAX_PATH];

	::ZeroMemory(&nr, sizeof(NETRESOURCE));
	::ZeroMemory(szTemp, MAX_PATH);
	_snprintf(szTemp, MAX_PATH, "%s\\%s", szServer, szShare);

	nr.dwType = RESOURCETYPE_ANY;
	nr.lpRemoteName = szTemp;
	nr.lpProvider = NULL;
	nr.lpLocalName = szDrive;
	nr.lpComment = "Added by fgdump";

	dwResult = WNetAddConnection2(&nr,
								  (LPSTR) NULL, 
								  (LPSTR) NULL,  
								  CONNECT_UPDATE_PROFILE);
	 
	if (dwResult == ERROR_ALREADY_ASSIGNED) 
	{ 
		Log.ReportError(CRITICAL, "Already connected to specified resource on drive %s.\n", szDrive); 
		return false; 
	} 
	else if (dwResult == ERROR_DEVICE_ALREADY_REMEMBERED) 
	{ 
		Log.ReportError(CRITICAL, "Attempted reassignment of remembered device %s.\n", szDrive); 
		return false; 
	} 
	else if(dwResult != NO_ERROR) 
	{ 
		Log.ReportError(CRITICAL, "A generic error occurred binding the drive: %d\n", dwResult); 
		return false; 
	} 
	 
	Log.ReportError(DEBUG, "The drive %s is bound to %s, testing write access...\n", szDrive, szShare);
	::ZeroMemory(szTempFilename, MAX_PATH);
	_snprintf(szTempFilename, MAX_PATH, "%s\\%s.fgdump", szDrive, szServer);

	std::ofstream outputFile(szTempFilename, std::ios::out | std::ios::trunc);
	outputFile.write("success", 7);
	if (outputFile.fail())
	{
		Log.ReportError(DEBUG, "Error writing the test file, skipping this share\n");
		WNetCancelConnection2(szDrive, CONNECT_UPDATE_PROFILE, TRUE);
		return false;
	}

	outputFile.flush();
	Log.ReportError(DEBUG, "Able to write to this directory, using drive %s for cachedump\n", szDrive);
	outputFile.close();
	DeleteFile(szTempFilename);

	return true;
}

void ShareFinder::UnbindDrive(char* szDrive)
{
	DWORD dwResult = WNetCancelConnection2(szDrive, CONNECT_UPDATE_PROFILE, TRUE);
	if (dwResult == ERROR_DEVICE_IN_USE)
	{
		Log.ReportError(CRITICAL, "Unable to disconnect drive %s, it appears another process is using the share. Suggest disconnecting it manually.\n", szDrive);
	}

	return;
}

char ShareFinder::GetUnusedDriveLetter()
{
	char szTemp[3] = { 0, 0, 0 };
	// Returns 0 if unsuccessful - that probably shouldn't happen too often!

	for (int i = 67; i <= 90; i++)	// 'C' through 'Z'
	{
		_snprintf(szTemp, 2, "%c:", i);
		if (GetDriveType(szTemp) == DRIVE_NO_ROOT_DIR)
		{
			// This drive should work
			Log.ReportError(DEBUG, "Drive %c is available, using that for bind operations\n", i);
			break;
		}
		szTemp[0] = 0;
	}

	return szTemp[0];
}
