#! /usr/bin/perl
# 
# $Id$
# Redirect system calls of the specified file system driver to TraceFS.sys
# Copyright (C) 2003 Jan Kratochvil <project-captive@jankratochvil.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; exactly version 2 of June 1991 is required
# 
# 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., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA


use strict;
use warnings;
use Carp qw(cluck confess);


my $D=0;

my $IMAGE_FILE_IMPORT_DIRECTORY=1;
my $IMAGE_FILE_EXPORT_DIRECTORY=0;
my $IMAGE_DATA_DIRECTORY_sizeof=8;
my $VirtualAddress_rel_to_DATA_DIRECTORY_offs=0;
my $Size_rel_to_DATA_DIRECTORY_offs=4;
my $IMAGE_EXPORT_DIRECTORY_sizeof=0x28;
my $IMAGE_IMPORT_DESCRIPTOR_sizeof=0x14;
my $Name_rel_to_IMAGE_EXPORT_DIRECTORY_offs=0xC;
my $NumberOfNames_rel_to_IMAGE_EXPORT_DIRECTORY_offs=0x18;
my $AddressOfNames_rel_to_IMAGE_EXPORT_DIRECTORY_offs=0x20;
my $OriginalFirstThunk_rel_to_IMPORT_DIRECTORY_offs=0x0;
my $FirstThunk_rel_to_IMPORT_DIRECTORY_offs=0x10;
my $Name_rel_to_IMAGE_IMPORT_BY_NAME_offs=0x2;
my $IMAGE_SECTION_HEADER_sizeof=0x28;
my $Name_rel_to_IMAGE_SECTION_HEADER_offs=0x00;
my $VirtualSize_rel_to_IMAGE_SECTION_HEADER_offs=0x08;
my $VirtualAddress_rel_to_IMAGE_SECTION_HEADER_offs=0x0C;
my $SizeOfRawData_rel_to_IMAGE_SECTION_HEADER_offs=0x10;
my $PointerToRawData_rel_to_IMAGE_SECTION_HEADER_offs=0x14;
my $Characteristics_rel_to_IMAGE_SECTION_HEADER_offs=0x24;


sub uintX_get($$$)
{
my($file,$offset,$bits)=@_;

	my $r=0;
	for (my $byte=0;$byte<$bits/8;$byte++) {
		confess if !defined $file->[$offset+$byte];
		$r|=($file->[$offset+$byte])<<($byte*8);
		}
	return $r;
}
sub uint32_get($$) { return uintX_get($_[0],$_[1],32); }
sub uint16_get($$) { return uintX_get($_[0],$_[1],16); }
sub  uint8_get($$) { return uintX_get($_[0],$_[1], 8); }

sub uintX_put($$$$)
{
my($file,$offset,$num,$bits)=@_;

	for (my $byte=0;$byte<$bits/8;$byte++) {
		confess if !defined $offset;
		$file->[$offset+$byte]=($num>>($byte*8))&0xFF;
		}
}
sub uint32_put($$$) { return uintX_put($_[0],$_[1],$_[2],32); }
sub uint16_put($$$) { return uintX_put($_[0],$_[1],$_[2],16); }
sub  uint8_put($$$) { return uintX_put($_[0],$_[1],$_[2], 8); }
sub uint32_push($$) { return uintX_put($_[0],@{$_[0]},$_[1],32); }
sub uint16_push($$) { return uintX_put($_[0],@{$_[0]},$_[1],16); }
sub  uint8_push($$) { return uintX_put($_[0],@{$_[0]},$_[1], 8); }


sub sum_offs($)
{
my($file)=@_;

	return uint32_get($file,0x3C)	# 3c: Offset to extended header
			+0x18	# OptionalHeader
			+0x40;	# CheckSum
}

sub sum_get($)
{
my($file)=@_;

	return uint32_get($file,sum_offs($file));
}

sub sum_put($$)
{
my($file,$sum)=@_;

	return uint32_put($file,sum_offs($file),$sum);
}

sub sum_calc($)
{
my($file)=@_;

	$file=[ @$file ];
	sum_put($file,0);
	my $length=@$file;
	my $sum=0;
	while (@$file) {
		$sum+=(shift @$file || 0) | (shift @$file || 0)<<8;
		$sum=($sum&0xFFFF)+($sum>>16);
		}
	$sum+=$length;
	return $sum;
}

sub align($$)
{
my($file,$align)=@_;

	push @$file,0 while @$file%$align;
}

sub stringz_put($$$)
{
my($file,$offset,$stringz)=@_;

	$stringz=[ map({ ord(); } split //,$stringz),0 ];
	while (@$stringz) {
		uint8_put($file,$offset,shift @$stringz);
		$offset++;
		}
}

sub stringz_push($$)
{
my($file,$stringz)=@_;

	stringz_put($file,@$file,$stringz);
	align $file,4;
}

sub stringz_get($$)
{
my($file,$offset)=@_;

	my $r="";
	while ((my $ord=uint8_get($file,$offset))) {
		$r.=chr $ord;
		$offset++;
		}
	return $r;
}

sub zeroes_push($$)
{
my($file,$count)=@_;

	push @$file,0 while $count-->0;
}

sub SizeOfImage_offs($)
{
my($file)=@_;

	return uint32_get($file,0x3C)	# 3c: Offset to extended header
				+0x18	# OptionalHeader
				+0x38;	# SizeOfImage
}

my $align=0x80;

sub sanity($$)
{
my($file,$file_id)=@_;

	my $calced=sum_calc($file);
	if ($calced!=sum_get($file)) {
		warn sprintf "$file_id: Original checksum wrong: found=0x%08X, calced=0x%08X",sum_get($file),$calced;
		}

	die sprintf "$file_id: Length 0x%X not aligned",scalar(@$file) if @$file%$align;
	my $SizeOfImage=uint32_get($file,SizeOfImage_offs($file));
	die sprintf "$file_id: SizeOfImage==0x%X but file size==0x%X",$SizeOfImage,scalar(@$file) if $SizeOfImage!=@$file;
}

sub export_names_get($)
{
my($tracefs)=@_;

	my $EXPORT_DIRECTORY_offs=uint32_get($tracefs,0x3C)	# 3c: Offset to extended header
				+0x18	# OptionalHeader
				+0x60	# DataDirectory
				+$IMAGE_FILE_EXPORT_DIRECTORY*$IMAGE_DATA_DIRECTORY_sizeof;
	my $EXPORT_DIRECTORY_VirtualAddress=uint32_get($tracefs,$EXPORT_DIRECTORY_offs+$VirtualAddress_rel_to_DATA_DIRECTORY_offs);
	my $EXPORT_DIRECTORY_Size=uint32_get($tracefs,$EXPORT_DIRECTORY_offs+$Size_rel_to_DATA_DIRECTORY_offs);
	die sprintf "EXPORT_DIRECTORY_Size 0x%X less than IMAGE_IMPORT_DESCRIPTOR_sizeof 0x%X",
			$EXPORT_DIRECTORY_Size,$IMAGE_IMPORT_DESCRIPTOR_sizeof if $EXPORT_DIRECTORY_Size<$IMAGE_IMPORT_DESCRIPTOR_sizeof;
	my $IMAGE_EXPORT_DIRECTORY_Name=uint32_get($tracefs,$EXPORT_DIRECTORY_VirtualAddress
			+$Name_rel_to_IMAGE_EXPORT_DIRECTORY_offs);
	my $tracefs_export_name=stringz_get($tracefs,$IMAGE_EXPORT_DIRECTORY_Name);
	my $tracefs_export_names_num=uint32_get($tracefs,$EXPORT_DIRECTORY_VirtualAddress
			+$NumberOfNames_rel_to_IMAGE_EXPORT_DIRECTORY_offs);
	my $tracefs_export_AddressOfNames=uint32_get($tracefs,$EXPORT_DIRECTORY_VirtualAddress
			+$AddressOfNames_rel_to_IMAGE_EXPORT_DIRECTORY_offs);
	my @tracefs_export_names;
	print STDERR "$tracefs_export_name exports:\n" if $D;
	for (my $namei=0;$namei<$tracefs_export_names_num;$namei++) {
		my $name_address=uint32_get($tracefs,$tracefs_export_AddressOfNames+4*$namei);
		my $name=stringz_get($tracefs,$name_address);
		push @tracefs_export_names,$name;
		print STDERR "\t$name\n" if $D;
		}
	
	return ($tracefs_export_name,@tracefs_export_names);
}


undef $/;
my $file=[ map({ ord(); } split //,<>) ];
my $tracefs=[ map({ ord(); } split //,<>) ];

sanity($tracefs,"tracefs");
my($tracefs_export_name,@tracefs_export_names)=export_names_get($tracefs);
for my $tname (@tracefs_export_names) {
	die "tracefs exported tname is not /^T/ compliant: $tname" if $tname!~/^T/;
	}
# FIXME: compiled to .sys tracefs contains internal export name .dll but we need import .sys !
$tracefs_export_name=~s/[.]dll$/.sys/i;

sanity($file,"file");


# import directory load

my $IMPORT_DIRECTORY_offs=uint32_get($file,0x3C)	# 3c: Offset to extended header
			+0x18	# OptionalHeader
			+0x60	# DataDirectory
			+$IMAGE_FILE_IMPORT_DIRECTORY*$IMAGE_DATA_DIRECTORY_sizeof;
my $IMPORT_DIRECTORY_VirtualAddress=uint32_get($file,$IMPORT_DIRECTORY_offs+$VirtualAddress_rel_to_DATA_DIRECTORY_offs);
my $IMPORT_DIRECTORY_Size=uint32_get($file,$IMPORT_DIRECTORY_offs+$Size_rel_to_DATA_DIRECTORY_offs);
die sprintf "IMPORT_DIRECTORY_Size 0x%X not aligned to IMAGE_IMPORT_DESCRIPTOR_sizeof 0x%X",
		$IMPORT_DIRECTORY_Size,$IMAGE_IMPORT_DESCRIPTOR_sizeof if $IMPORT_DIRECTORY_Size%$IMAGE_IMPORT_DESCRIPTOR_sizeof;

my $IMPORT_DIRECTORY=[ @{$file}[$IMPORT_DIRECTORY_VirtualAddress
		..($IMPORT_DIRECTORY_VirtualAddress+$IMPORT_DIRECTORY_Size-1)] ];
uintX_put($file,$IMPORT_DIRECTORY_VirtualAddress,0,8*$IMPORT_DIRECTORY_Size);	# zero the original space
for (my $zerotail=0;$zerotail<$IMAGE_IMPORT_DESCRIPTOR_sizeof;$zerotail++) {
	die "IMPORT_DIRECTORY tail not zeroed" if pop @$IMPORT_DIRECTORY;
	}


# import directory entries processing

my %name_to_FirstThunk;	# name->FirstThunk
my %tname_to_name;	# string->string{^.->T}
for (
		my $IMPORT_DIRECTORY_offset=0;
		$IMPORT_DIRECTORY_offset<@$IMPORT_DIRECTORY;
		$IMPORT_DIRECTORY_offset+=$IMAGE_IMPORT_DESCRIPTOR_sizeof) {
	my $OriginalFirstThunk_base=uint32_get($IMPORT_DIRECTORY,$IMPORT_DIRECTORY_offset
			+$OriginalFirstThunk_rel_to_IMPORT_DIRECTORY_offs);
	my $FirstThunk_base=uint32_get($IMPORT_DIRECTORY,$IMPORT_DIRECTORY_offset
			+$FirstThunk_rel_to_IMPORT_DIRECTORY_offs);
	for (my $OriginalFirstThunk=$OriginalFirstThunk_base;;$OriginalFirstThunk+=4) {
		my $AddressOfData=uint32_get($file,$OriginalFirstThunk);
		last if !$AddressOfData;
		my $name=stringz_get($file,$AddressOfData+$Name_rel_to_IMAGE_IMPORT_BY_NAME_offs);
		print STDERR "import $name\n" if $D;
		die "Invalid name import as it has leading 'T': $name" if $name=~/^T/;
		(my $tname=$name)=~s/^./T/;
		die "Name conflict in tname map: $name->$tname" if exists $tname_to_name{$tname};
		$tname_to_name{$tname}=$name;
		die if exists $name_to_FirstThunk{$name};
		$name_to_FirstThunk{$name}=$FirstThunk_base+($OriginalFirstThunk-$OriginalFirstThunk_base);
		}
	}


# add-on import directories generation
my $addon_section_base=@$file;

my $tracefs_export_name_offs=@$file;
stringz_push($file,$tracefs_export_name);

my $ordinal=1;
for my $tname (@tracefs_export_names) {
	die "tracefs exported tname $tname not found in imports" if !exists $tname_to_name{$tname};
	my $name=$tname_to_name{$tname};
	my $FirstThunk=$name_to_FirstThunk{$name};

	my $IMAGE_IMPORT_BY_NAME_offs=@$file;
	uint16_push($file,$ordinal++);	# Hint
	stringz_push($file,$tname);	# Name

	my $OriginalFirstThunk_offs=@$file;
	uint32_push($file,$IMAGE_IMPORT_BY_NAME_offs);
	uint32_push($file,0);	# zero terminator

	uint32_push($IMPORT_DIRECTORY,$OriginalFirstThunk_offs);
	uint32_push($IMPORT_DIRECTORY,0);	# TimeDateStamp
	uint32_push($IMPORT_DIRECTORY,0);	# ForwarderChain
	uint32_push($IMPORT_DIRECTORY,$tracefs_export_name_offs);	# Name
	uint32_push($IMPORT_DIRECTORY,$FirstThunk);
	}

zeroes_push $IMPORT_DIRECTORY,$IMAGE_IMPORT_DESCRIPTOR_sizeof;
uint32_put($file,$IMPORT_DIRECTORY_offs+$VirtualAddress_rel_to_DATA_DIRECTORY_offs,scalar(@$file));
uint32_put($file,$IMPORT_DIRECTORY_offs+$Size_rel_to_DATA_DIRECTORY_offs,scalar(@$IMPORT_DIRECTORY));
push @$file,@$IMPORT_DIRECTORY;


# rebuild the sections for add-on data
align $file,$align;
# concatenate .rdata to .data
my $SizeOfOptionalHeader_offs=uint32_get($file,0x3C)	# 3c: Offset to extended header
			+0x04	# FileHeader
			+0x10;	# SizeOfOptionalHeader
my $SizeOfOptionalHeader=uint16_get($file,$SizeOfOptionalHeader_offs);
my $NumberOfSections_offs=uint32_get($file,0x3C)	# 3c: Offset to extended header
			+0x04	# FileHeader
			+0x2;	# NumberOfSections
my $NumberOfSections=uint16_get($file,$NumberOfSections_offs);
my $IMAGE_SECTION_HEADER_offs=uint32_get($file,0x3C)	# 3c: Offset to extended header
			+0x18	# OptionalHeader
			+$SizeOfOptionalHeader;

my($offset_data,$offset_rdata);
for (
		my $IMAGE_SECTION_HEADER_offset=$IMAGE_SECTION_HEADER_offs;
		$IMAGE_SECTION_HEADER_offset<$IMAGE_SECTION_HEADER_offs+$NumberOfSections*$IMAGE_SECTION_HEADER_sizeof;
		$IMAGE_SECTION_HEADER_offset+=$IMAGE_SECTION_HEADER_sizeof) {
	my $VirtualAddress=uint32_get($file,$IMAGE_SECTION_HEADER_offset+$VirtualAddress_rel_to_IMAGE_SECTION_HEADER_offs);
	my $PointerToRawData=uint32_get($file,$IMAGE_SECTION_HEADER_offset+$PointerToRawData_rel_to_IMAGE_SECTION_HEADER_offs);
	die sprintf "VirtualAddress 0x%X != PointerToRawData 0x%X in IMAGE_SECTION_HEADER at 0x%X",
			$VirtualAddress,$PointerToRawData,$IMAGE_SECTION_HEADER_offset if $VirtualAddress!=$PointerToRawData;
	my $SizeOfRawData=uint32_get($file,$IMAGE_SECTION_HEADER_offset+$SizeOfRawData_rel_to_IMAGE_SECTION_HEADER_offs);
	my $VirtualSize=uint32_get($file,$IMAGE_SECTION_HEADER_offset+$VirtualSize_rel_to_IMAGE_SECTION_HEADER_offs);
	$VirtualSize+=$align-1;
	$VirtualSize-=$VirtualSize%$align;
	die sprintf "up_align(VirtualSize,0x%X) 0x%X != SizeOfRawData 0x%X in IMAGE_SECTION_HEADER at 0x%X",
			$align,$VirtualSize.$SizeOfRawData,$IMAGE_SECTION_HEADER_offset if $VirtualSize!=$SizeOfRawData;
	my $Characteristics=uint32_get($file,$IMAGE_SECTION_HEADER_offset+$Characteristics_rel_to_IMAGE_SECTION_HEADER_offs);
	my $is_data =($Characteristics==0xC8000040);
	my $is_rdata=($Characteristics==0x48000040);
	die sprintf "Duplicate .data in IMAGE_SECTION_HEADER at 0x%X",$IMAGE_SECTION_HEADER_offset if $is_data && $offset_data;
	die sprintf "Duplicate .rdata in IMAGE_SECTION_HEADER at 0x%X",$IMAGE_SECTION_HEADER_offset if $is_rdata && $offset_rdata;
	$offset_data=$IMAGE_SECTION_HEADER_offset if $is_data;
	$offset_rdata=$IMAGE_SECTION_HEADER_offset if $is_rdata;
	}
die ".data section not found" if !$offset_data;
die ".rdata section not found" if !$offset_rdata;
my $data_PointerToRawData=uint32_get($file,$offset_data+$PointerToRawData_rel_to_IMAGE_SECTION_HEADER_offs);
my $data_SizeOfRawData=uint32_get($file,$offset_data+$SizeOfRawData_rel_to_IMAGE_SECTION_HEADER_offs);
my $data_VirtualSize=uint32_get($file,$offset_data+$VirtualSize_rel_to_IMAGE_SECTION_HEADER_offs);
my $rdata_PointerToRawData=uint32_get($file,$offset_rdata+$PointerToRawData_rel_to_IMAGE_SECTION_HEADER_offs);
my $rdata_SizeOfRawData=uint32_get($file,$offset_rdata+$SizeOfRawData_rel_to_IMAGE_SECTION_HEADER_offs);
my $rdata_VirtualSize=uint32_get($file,$offset_rdata+$VirtualSize_rel_to_IMAGE_SECTION_HEADER_offs);
die ".data is not right after .rdata" if $rdata_PointerToRawData+$rdata_SizeOfRawData!=$data_PointerToRawData;
uint32_put($file,$offset_data+$PointerToRawData_rel_to_IMAGE_SECTION_HEADER_offs,$rdata_PointerToRawData);
uint32_put($file,$offset_data+$VirtualAddress_rel_to_IMAGE_SECTION_HEADER_offs,$rdata_PointerToRawData);
uint32_put($file,$offset_data+$SizeOfRawData_rel_to_IMAGE_SECTION_HEADER_offs,$rdata_SizeOfRawData+$data_SizeOfRawData);
uint32_put($file,$offset_data+$VirtualSize_rel_to_IMAGE_SECTION_HEADER_offs,$rdata_SizeOfRawData+$data_VirtualSize);
# .rdata coalesced to .data, .rdata to be rebuilt now:
uintX_put($file,$offset_rdata,0,8*$IMAGE_SECTION_HEADER_sizeof);	# zero the original space
stringz_put($file,$offset_rdata+$Name_rel_to_IMAGE_SECTION_HEADER_offs,"INIT");
uint32_put($file,$offset_rdata+$VirtualSize_rel_to_IMAGE_SECTION_HEADER_offs,@$file-$addon_section_base);
uint32_put($file,$offset_rdata+$VirtualAddress_rel_to_IMAGE_SECTION_HEADER_offs,$addon_section_base);
uint32_put($file,$offset_rdata+$SizeOfRawData_rel_to_IMAGE_SECTION_HEADER_offs,@$file-$addon_section_base);
uint32_put($file,$offset_rdata+$PointerToRawData_rel_to_IMAGE_SECTION_HEADER_offs,$addon_section_base);
uint32_put($file,$offset_rdata+$Characteristics_rel_to_IMAGE_SECTION_HEADER_offs,0xE2000020);


# file output finalization
uint32_put($file,SizeOfImage_offs($file),scalar(@$file));
sum_put($file,sum_calc($file));
print join("",map({ chr; } @$file));
