#include <stdio.h>
#include "peread.h"
#include <sys/stat.h>
#include <fcntl.h>
#include <stdlib.h>
#include <string.h>
char* fileOpen(char *path);
int fileSave(char*, struct IMAGE_DOS_HEADER*,struct IMAGE_NT_HEADERS*, struct IMAGE_SECTION_HEADER**, char **);
struct IMAGE_SECTION_HEADER* loadSection(struct IMAGE_DOS_HEADER*,struct IMAGE_NT_HEADERS *,char *,int);
int
main(int argc,char **argv){
	if(argc != 3){
		printf("%s : missing file operand\n",argv[0]);
		return -1;
	}
	int i;
	struct IMAGE_DOS_HEADER *d_hdr;
	struct IMAGE_NT_HEADERS *nt_hdr;
	struct IMAGE_SECTION_HEADER *s_hdr;
	char *file_data;
	file_data = fileOpen(argv[1]);
	if( file_data == NULL){
		return 1;
	}
	d_hdr = (struct IMAGE_DOS_HEADER *) file_data;
	nt_hdr = (struct IMAGE_NT_HEADERS *) (file_data + d_hdr->e_lfanew);
	s_hdr = (struct IMAGE_SECTION_HEADER *) (file_data+d_hdr->e_lfanew+sizeof(struct IMAGE_NT_HEADERS));
	struct IMAGE_SECTION_HEADER *secs[nt_hdr->e_numsec];
	char*  sections[nt_hdr->e_numsec];
	char* pe_hdr;
	pe_hdr = malloc(nt_hdr->e_opthdr.e_secalgn);
	memcpy(pe_hdr,file_data,1024);
	for(i = 0; i < nt_hdr->e_numsec; i++)
	{
		secs[i] = loadSection(d_hdr,nt_hdr,file_data,i);
		printf("SECTION NAME: %s\nVirtual Size: 0x%x\nVirtual Address: 0x%x\nSize Of Raw data: 0x%x\n\n",
				secs[i]->s_name,secs[i]->s_virtsize,secs[i]->s_virtaddr,secs[i]->s_rwdtptr);
	}
	for(i = 0; i < nt_hdr->e_numsec; i++)
        {
		int section_size = ((secs[i]->s_virtaddr + nt_hdr->e_opthdr.e_secalgn -1)
					/nt_hdr->e_opthdr.e_secalgn)*nt_hdr->e_opthdr.e_secalgn ;
//		printf("0x%x %x\n",section_size, secs[i]->s_rwdtsize);
                sections[i] = malloc(section_size);
		memcpy(sections[i],(char *) (file_data+secs[i]->s_rwdtptr),section_size);
		printf("Copied 0x%x bytes to %s \n",section_size,secs[i]->s_name);
//		int j;
//		if (i == 0)
//		for (j = 0; j < section_size;){
//			printf("%2.2x",(unsigned char)sections[i][j++]);
//
//		}
//		printf("%.2x\n",sections[i][1]);
        }
	fileSave(argv[2],d_hdr,nt_hdr,secs,sections);
}


struct IMAGE_SECTION_HEADER *
loadSection(struct IMAGE_DOS_HEADER *d_hdr,struct IMAGE_NT_HEADERS *nt_hdr,char* importedFile,int secoff){

	struct IMAGE_SECTION_HEADER *s_hdr;
	s_hdr = (struct IMAGE_SECTION_HEADER *) (importedFile + d_hdr->e_lfanew + sizeof(struct IMAGE_NT_HEADERS) +
									(sizeof(struct IMAGE_SECTION_HEADER)*secoff) );
//	printf("%s\n",s_hdr->s_name);
	return s_hdr;


}

int
fileSave(char *path,struct IMAGE_DOS_HEADER *d_hdr, struct IMAGE_NT_HEADERS *nt_hdr,
		struct IMAGE_SECTION_HEADER **sec_hdr,char** sections){
	struct stat *st;
	char *out;
	int fd,wf,written,i;
	u_int32_t fileSize;
	fileSize = sec_hdr[nt_hdr->e_numsec -1]->s_rwdtptr + sec_hdr[nt_hdr->e_numsec -1]->s_rwdtsize;
	printf("%x\n",fileSize);
	out = malloc (fileSize);
	//memset(out,'\0',1000);
	memcpy(out,(char *) d_hdr,sizeof(struct IMAGE_DOS_HEADER));
	written = sizeof(struct IMAGE_DOS_HEADER);
//	memcpy(out+written,((char *)d_hdr)+written,0xb0);
//	written += 0xb0;

	if(written < d_hdr->e_lfanew){
		int n = d_hdr->e_lfanew - written;
		printf("writing at offset %x\n",n);
		memcpy(out+written, ((char *) (d_hdr))+written,n);
		written += n;
		printf("written : %x\n",written);
	}
	memcpy(out+written, (char *) nt_hdr, sizeof(struct IMAGE_NT_HEADERS));
	written += sizeof(struct IMAGE_NT_HEADERS);
	for(i=0; i < nt_hdr->e_numsec; i++){
		memcpy(out+written,sec_hdr[i],sizeof(struct IMAGE_SECTION_HEADER));
		written += sizeof(struct IMAGE_SECTION_HEADER);

	}
	for(i=0; i< nt_hdr->e_numsec; i++){
                int section_size = ((sec_hdr[i]->s_virtsize + nt_hdr->e_opthdr.e_secalgn -1)
                                        /nt_hdr->e_opthdr.e_secalgn)*nt_hdr->e_opthdr.e_secalgn;
		memcpy(out+sec_hdr[i]->s_rwdtptr,sections[i],section_size);
		printf("writing section %d with size %x\n",i,section_size);

	}
	// *** OPEN FILE *** //
	if( (fd = open(path,O_WRONLY | O_CREAT)) < 0){
		printf("error: Could not access file\n");
		return -1;
	}
	if( (wf=write(fd,out,fileSize)) < 0){
		printf("error: Could not write file\n");
		return -1;
	}
	close(fd);
	return 	0;
}
char*
fileOpen(char *path){
	int fd;
	struct stat st;
	char *in;
	fd = open(path,O_RDONLY);
	if(fd < 0){
		printf("error: could not read from file\n");
		return NULL;
	}
	if(fstat(fd,&st) < 0){
		printf("error: could not access file\n");
		return NULL;
	}

	in = malloc(st.st_size+1);
	if(in == NULL){
		printf("error: could not allocate memory\n");
		return NULL;
	}
	if( read(fd,in,st.st_size) < 0){
		printf("error: could not read file\n");
		return NULL;
	}
	in[st.st_size] = 0;
	close(fd);
	return in;
}