/*
	SelfExtract Class - attaching/detaching file to an executable.
	Written by Nir Dremer
*/

#include "SelfExtract.h"
#include <stdio.h>
#include <fcntl.h>
#include <Shlwapi.h>

namespace GenLib
{

SelfExtract::SelfExtract(const char *filename)
{
	// if no filename given, using current module filename.
	if (!filename) 
	{
		GetModuleFileName(0, _sourceFile, MAX_PATH);
	}
	else
	{
		strcpy(_sourceFile, filename);
	}
	_lastDetached[0] = 0;
}

SelfExtract::~SelfExtract(void)
{
}

int SelfExtract::attachFile(const char *filename, const char *writeName)
{
	// parameter must be valid.
	if (!filename) return 1;
	// file (given in constructor) must not have attached file!
	if (checkSignature()) return 2;

	FILE *orig = fopen(_sourceFile, "ab");
	if (orig == 0) return 3;

	FILE *attachy = fopen(filename, "rb");
	if (attachy == 0) 
	{
		fclose(orig);
		return 4;
	}
	
	fseek(orig, 0, SEEK_END);

	// save file beginning position.
	long offset = ftell(orig);

	const char *shortName = filename;
   	if (writeName != 0) shortName = writeName;
	// stripping directory information - only filename will be saved.
	if (strrchr(shortName, '\\')) shortName = strrchr(shortName, '\\')+1;
	
	int	fileLen = (int)strlen(shortName);
	
	// writing filename length (integer size) and filename.
	fwrite(&fileLen, sizeof(char), sizeof(fileLen), orig);
	fwrite(shortName, sizeof(char), fileLen, orig);

	const int bufSize = 1024;
	char *buffer[bufSize];
	size_t writeSize;
	
	while (!feof(attachy))
	{
		writeSize = fread(buffer, sizeof(char), bufSize, attachy);
		fwrite(buffer, sizeof(char), writeSize, orig);
	}

    fwrite(&offset, sizeof(char), sizeof(offset), orig);
	fwrite(ATT_SIGNATURE, sizeof(char), ATT_SIGSIZE, orig);

	fclose(orig);
	fclose(attachy);
	return true;
}

int SelfExtract::detachFile(const char *writeFile, char *detachedFile, int detachedSize)
{
	FILE *orig = fopen(_sourceFile, "rb");
	if (orig == 0) return 2;

	// seeking signature.
	if (fseek(orig, - (long) (ATT_SIGSIZE + sizeof(long)), SEEK_END) != 0)
	{
		fclose(orig);
		return 3;
	}

	long fileEndPos = ftell(orig);

	// reading the position of the beginning of filename + file content.
	long filePos;
	fread(&filePos, sizeof(char), sizeof(filePos), orig);

	// reading signature.
	char sig[ATT_SIGSIZE+1];
	sig[ATT_SIGSIZE] = 0;
	fread(sig, sizeof(char), ATT_SIGSIZE, orig);

	// signature check failure.
	if (strcmp(sig, ATT_SIGNATURE) != 0)
	{
		fclose(orig);
		return 3;
	}

	// going to file content location.
	if (fseek(orig, filePos, SEEK_SET) != 0)
	{
		fclose(orig);
		return 4;
	}

	// reading filename (and filename length).
	int nameLen;
	fread(&nameLen, sizeof(char), sizeof(nameLen), orig);
	char filename[MAX_PATH];
	fread(filename, sizeof(char), nameLen, orig);
	filename[nameLen] = 0;
	// Saving filename for user query - getDetached().
	strcpy(_lastDetached, filename);
 
	// building output file based on given directory and extracted filename.
	char outFile[MAX_PATH];
	// if there's a given folder.
	if (writeFile)
	{
		strcpy(outFile, writeFile);
	}
	else
	{
		strcpy(outFile, filename);
	}
	
	// copying the output file to the out parameter if param. valid.
	if (detachedFile)
		strncpy(detachedFile, outFile, detachedSize);

	FILE *attachy = fopen(outFile, "wb");

	// many reasons for fopen to fail:
	//	* result directory can NOT be writen (premission or protected media).
	//	* bad filename inside the executable (currently not begin checked).
	//	* and probably more reasons.
	if (attachy == 0) 
	{
		fclose(orig);
		return 5;
	}

	// open out file.
	const int bufSize = 1024;
	char *buffer[bufSize];
	int readSize = bufSize;

	filePos = ftell(orig);

	while (filePos < fileEndPos)
	{
		if (filePos + bufSize > fileEndPos)
		{
			readSize = fileEndPos - filePos;
		}
		
		filePos += (long)fread(buffer, sizeof(char), readSize, orig);
		fwrite(buffer, sizeof(char), readSize, attachy);
	}

	fclose(orig);
	fclose(attachy);
	return true;
}

bool SelfExtract::checkSignature()
{
	FILE *file = fopen(_sourceFile, "rb");
	if (file == 0) return false;

	if (fseek(file, - (int) ATT_SIGSIZE, SEEK_END) != 0)
	{
		fclose(file);
		return false;
	}

	// checking signature.
	char sig[ATT_SIGSIZE+1];
	sig[ATT_SIGSIZE] = 0;
	fread(sig, sizeof(char), ATT_SIGSIZE, file);

	// signature check failure.
	if (strcmp(sig, ATT_SIGNATURE) != 0)
	{
		fclose(file);
		return false;
	}

	fclose(file);
	return true;
}

}