#!/usr/bin/perl
# Extract all of the files from a UEFI firmware image.
#
# This is a simple replacement for uefi-firmware-parser
# with fewer features and more suited for reconstituting the
# firmware later.
#
use warnings;
use strict;
use FindBin;
use lib "$FindBin::Bin/../lib";
use EFI;
use Getopt::Long;
use File::Basename;
use File::Temp 'tempfile';
use Data::Dumper;

my $usage = <<"END";
Usage:
  $0 [options] firmware.rom | tee firmware.txt
Options:
  -h | -? | --help        this usage
  -v | --version          Print per-section logs
  -o | --output-dir DIR   Write the output to a different base directory
  -s | --start 0xX        Starting offset for partial extraction
  -n | --length 0xX       Number of bytes to extract for partial extraction
  -r | --repack           Unpack LZMA compressed file sections

END

my $start_offset = 0; # 0x02c00000;
my $extract_length;
my $verbose = 0;
my $base_dir = '.';
my $repack;

GetOptions(
	"h|?|help"		=> sub { print $usage; exit 0 },
	"v|verbose+"		=> \$verbose,
	"o|output-dir=s"	=> \$base_dir,
	"s|start=o"		=> \$start_offset,
	"n|length=o"		=> \$extract_length,
	"r|repack+"		=> \$repack,
) or die $usage;

local $/;
while(<>)
{
	process_region($base_dir, $_);
}

sub process_region
{
	my $base = shift;
	my $data = shift;
	my $length = length($data);

	printf "%s: length 0x%x\n", $base, $length
		if $verbose;

	my $start_unknown;

	# Search for the start of firmware volumes,
	# identified by their '_FVH' in the structure
	my $step = 256;

	# Adjust the end offset if they ask for too much
	my $end_offset = $length - 0x30;
	if (defined $extract_length)
	{
		my $new_end_offset = $start_offset + $extract_length;
		$end_offset = $new_end_offset
			if $new_end_offset < $end_offset;
	}

	for(my $offset = $start_offset ; $offset < $end_offset ; $offset += $step)
	{
		# Look for a flash region descriptor
		my $ifd_sig = unpack("N", substr($_, $offset + 0x10, 4));
		if ($ifd_sig == 0x5AA5F00F)
		{
			printf "%s/0x%08x.ifd: Flash Descriptor\n",
				$base,
				$offset,
				;

			my $ifd_len = 0x10000;
			my $data = substr($_, $offset, $ifd_len);

			output(sprintf("%s/0x%08x.ifd", $base, $offset), $data);

			$offset += $ifd_len - $step;
			next;
		}

		my $fv_sig = substr($_, $offset + 0x28, 4);
		my $fv_length = EFI::read64($_, $offset + 0x20);

		if ($fv_sig ne '_FVH' or $offset + $fv_length > $length)
		{
			# likely not a filesystem; report an unknown region
			# if we have started processing filesystems
			$start_unknown = $offset
				unless defined $start_unknown;
			next;
		}

		if (defined $start_unknown)
		{
			# We have a unknown region to write out
			my $len = $offset - $start_unknown;
			my $data = substr($_, $start_unknown, $len);
			printf "%s/0x%08x.bin: UNKNOWN length 0x%x \n",
				$base, $start_unknown, $len;

			output(sprintf("%s/0x%08x.bin", $base, $start_unknown), $data, 1);
			undef $start_unknown;
		}

		my $fv = substr($data, $offset, $fv_length);

		process_fv(sprintf("%s/0x%08x", $base, $offset), $fv);

		# skip to the end of the filesystem
		# should we care if this FV is not processed?
		$offset += $fv_length - $step;
	}
}


sub process_fv
{
	my $base = shift;
	my $fv = shift;
	my $fv_length = length $fv;
	my $guid = EFI::read_guid($fv, 16);

	output("$base.fv", $fv);
	printf "%s.fv: FV %s length 0x%x\n",
		$base,
		$guid,
		$fv_length,
		;

	if ($guid ne '8c8ce578-8a3d-4f1c-9935-896185c32dd3'
	and $guid ne '5473c07a-3dcb-4dca-bd6f-1e9689e7349a'
	) {
		# we can only process normal firmware volumes
		return;
	}

	# read the start of data offset from the header
	my $offset = EFI::read16($fv, 0x30);
	if ($offset >= $fv_length)
	{
		die sprintf "%s: FV invalid data offset 0x%04x\n",
			$base, $offset;
	}


	while($offset < $fv_length - 0x20)
	{
		my $len = EFI::read24($fv, $offset + 0x14);
		my $data_offset = 0x18;

		if ($len == 0xFFFFFF)
		{
			# Version 2 header with extended length
			$len = EFI::read64($fv, $offset + 0x18);

			# If we have an all-0xFF length, which indicates
			# the start of free space
			return 1 if ~$len == 0;

			## Looks good, adjust the starting offset
			$data_offset += 0x8;
		}

		if ($len == 0x0)
		{
			die sprintf "%s: 0x%08x file has zero length?\n",
				$base,
				$offset;
		}

		if ($len + $offset > $fv_length)
		{
			warn sprintf "%s: 0x%08x file len 0x%x exceeds FV len\n",
				$base,
				$offset,
				$len;
			return;
		}


		my $data = substr($fv, $offset, $len);

		process_ffs($base, $data, $data_offset);

		$offset += $len;

		# align it
		$offset = ($offset + 7) & ~7;
	}

	return 1;
}


sub process_ffs
{
	my $base = shift;
	my $ffs = shift;
	my $data_offset = shift; # might be 0x18 or 0x20
	my $len = length($ffs);

	my $guid = EFI::read_guid($ffs, 0x00);
	my $type = ord(substr($ffs, 0x12, 1));

	if ($guid eq 'ffffffff-ffff-ffff-ffff-ffffffffffff')
	{
		# padding: do not output it
		# should check that everything is 0xFF first
		return;
	}

	output("$base/$guid.ffs", $ffs);
	my $data = substr($ffs, $data_offset);
	my $type_str = EFI::file_type_lookup($type);
	my $name = $type_str;

	if ($name eq 'FFS_PAD' or $name eq 'RAW')
	{
		# we are done here
		printf "%s/%s.ffs: $name length 0x%x \n", $base, $guid, $len;
		return;
	}

	# if we are re-packing, recursively expand the sections
	# and regenerate the section file
	my @sections = process_sections("$base/$guid", $data);
	my @new_sections = map { EFI::section($_->[0], $_->[1]) } @sections;

	if ($repack)
	{
		output("$base/$guid.ffs", EFI::ffs(
			$type_str,
			EFI::guid($guid),
			@new_sections
		));
	}

	# if this file has a UI section, extract its name
	$name = EFI::read_ucs16($_->[1], 0)
		for grep { $_->[0] eq 'USER_INTERFACE'} @sections;

	print "$base/$guid.ffs: $name\n";

}


sub process_sections
{
	my $base = shift;
	my $data = shift;
	my @rc;

	for (ffs_sections($base, $data))
	{
		my @sec = process_section($base, @$_);
		push @rc, @sec if @sec;
	}

	return @rc;
}

sub process_section
{
	my $base = shift;
	my $type = shift;
	my $sec = shift;
	my $number = shift;
	my $offset = shift;

	if ($type eq 'USER_INTERFACE')
	{
		# we should record that we've seen a UI section
		#printf "Name: %s\n", EFI::read_ucs16($sec, 0);
	}

	# Look for GUID defined sections that encode LZMA compressed data
	if ($type eq 'GUID_DEFINED')
	{
		my $guid = EFI::read_guid($sec, 0);
		return unless $guid eq $EFI::lzma_guid;

		my $lz_data = substr($sec, 0x14, length($sec) - 0x14);

		my ($fh,$filename) = tempfile();
		print $fh $lz_data;
		close $fh;

		my $data = `lzma --decompress --stdout - < $filename`;
		printf "%s.ffs: length 0x%x (0x%x compressed)\n",
			$base,
			length($data),
			length($lz_data),
			;

		# recursively process this data
		return process_sections($base, $data);
	}

	if ($type eq 'FIRMWARE_VOLUME_IMAGE')
	{
		process_fv(sprintf("%s/%d", $base, $number), $sec);
		return;
	}

	if ($type eq 'RAW'
	or  $type eq 'TIANO_COMPRESSED'
	) {
		#return;
	}

	warn "$base/$number: $type unknown\n"
		if $type =~ /^0x/;

	return ([ $type, $sec, $number, $offset ]);
}


sub ffs_sections
{
	my $base = shift;
	my $ffs = shift;
	my $ffs_len = length($ffs);

	my @sections;

	# find each section inside
	my $number = 0;
	my $offset = 0;
	while($offset < $ffs_len - 8)
	{
		my $len = EFI::read24($ffs, $offset);
		my $data_offset = 0x4;
		if ($len == 0xFFFFFF)
		{
			# FFSv3 section
			$len = EFI::read32($ffs, $offset + $data_offset);
			$data_offset += 4;
		}

		if ($len < $data_offset)
		{
			warn sprintf "%s: 0x%x Section length %x invalid\n", $base, $offset, $len;
			return;
		}

		if ($offset + $len > $ffs_len)
		{
			die sprintf "%s: Section length %x exceeds FFS len %x\n",
				$base,
				$len,
				$ffs_len,
				;
		}

		my $sec = substr($ffs, $offset, $len);

		# move to the next section, keeping a 4-byte alignment
		$offset = ($offset + $len + 3) & ~3;

		my $sec_type = ord(substr($sec, 3, 1));
		my $sec_data = substr($sec, $data_offset, $len - $data_offset);

		my $sec_type_name = $EFI::section_types_lookup{$sec_type};
		$sec_type_name ||= sprintf "0x%02x", $sec_type;

		printf "%s.ffs: %s len %x\n", $base, $sec_type_name, $len
			if $verbose;

		# ignore RAW sections that are all padding
		next if $sec_type_name eq 'RAW' and empty_data($sec_data);

		push @sections, [ $sec_type_name, $sec_data, $number, $offset ];
		$number++;
	}

	return @sections;
}


sub empty_data
{
	my $data = shift;
	my $byte = substr($data, 0, 1);
	for(my $i = 0 ; $i < length $data ; $i++)
	{
		return 0 if substr($data, $i, 1) ne $byte;
	}

	return 1;
}


sub output
{
	my $name = shift,
	my $data = shift;
	my $force = shift;

	# check for an empty region (all the same value)
	return if not $force and empty_data($data);

	my $dir = dirname($name);
	system(mkdir => -p => $dir)
		and die "$name: Unable to create directory\n";
	open FILE, '>', "$name"
		or die "$name: Unable to create output file: $!\n";
	print FILE $data;
	close FILE;
}