#!/usr/bin/perl
our $VERSION = "0.4";

# fitimg - (C) 2021 Ralf Steines aka Hippie2000 - <metamonk@yahoo.com>
# Handle and manipulate firmware images in AVM /var/tmp/fit-image format.
# Docs and latest version can be found at https://boxmatrix.info/wiki/FIT-Image

# 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, see <http://www.gnu.org/licenses/>.

use strict;
use warnings; # requires perl 5.6
use Getopt::Std; # Perl core module (part of Perl package)
use String::CRC32; # May need 'cpan install String::CRC32'

# Freetz filenames
my %fname = (
	# HWR 257 - 5530
	257 => {
		'prxI_HW257_kernel' => 'kernel.image',
		'prxI_HW257_ramdisk' => 'filesystem.image',
		'prxI_HW257_flat_dt_0_aon' => 'flatdt_0_aon.image',
		'prxI_HW257_flat_dt_0_pon' => 'flatdt_0_pon.image',
		'prxB_HW0257_kernel' => 'kernel2.image',
		'prxB_HW0257_ramdisk' => 'filesystem2.image',
		'prxB_HW0257_flat_dt_0' => 'flatdt2_0.image',
	},
	# HWR 256 - 7530ax
	256 => {
		'brcma9_HW256_kernel' => 'kernel.image',
		'brcma9_HW256_squashFS_filesystem' => 'filesystem.image',
		'brcma9_HW256_flat_dt_1' => 'flatdt_1.image',
		'brcma9_HW256_flat_dt_0' => 'flatdt_0.image',
		'brcma9TZ_HW256_kernel' => 'kernel2.image',
		'brcma9TZ_HW256_flat_dt_0' => 'flatdt2_0.image',
	},
	# HWR 253 - 6000
	253 => {
		'qcaarmv8_HW253_kernel' => 'kernel.image', 
		'qcaarmv8_HW253_squashFS_filesystem' => 'filesystem.image',
		'qcaarmv8_HW253_flat_dt_0' => 'flatdt_0.image',
		'qcaarmv8_HW253_flat_dt_2' => 'flatdt_2.image',
	},
);
my %missing = %fname;

# string hunk names
our %strname = (
	# all models
	0 => {
		0 => 'info',
		65 => 'data', # blob
		70 => 'type',
		75 => 'arch',
		80 => 'os',
		83 => 'compression',
		106 => 'hashalgo',
	},
	# HWR 257 - 5530
	257 => {
		117 => 'avm,endianess',
		252 => 'kernel',
		259 => 'flatdt',
		263 => 'ramdisk',
	},
	# HWR 256 - 7530ax
	256 => {
		117 => 'avm,kernel-args',
		133 => 'avm,endianess',
		283 => 'kernel',
		290 => 'flatdt',
		294 => 'squashfs',
	},
	# HWR 253 - 6000
	253 => {
		117 => 'avm,endianess',
		267 => 'avm,kernel-args',
		283 => 'kernel',
		290 => 'flatdt',
		294 => 'squashfs',
	},
);

# 32-bit word hunk names
# FIXME: some few are shorter than 32bit
our %wordname = (
	# global
	0 => {
		12 => 'timestamp',
		22 => 'avm,gu-version',
		37 => 'avm,variants',
		50 => '#address-cells',
		95 => 'loadaddr',
		100 => 'entryaddr',
		111 => 'hashvalue',
	},
	# HWR 257 - 5530
	257 => {
		131 => 'avm,kernel_text_start',
		153 => 'avm,names',
		163 => 'avm,token_table',
		179 => 'avm,token_index',
		195 => 'avm,num_syms',
		208 => 'avm,addresses',
		222 => 'avm,relative_base',
		240 => 'avm,offsets',
	},
	# HWR 256 - 7530ax
	256 => {
		147 => 'avm,kernel_text_start',
		169 => 'avm,names',
		179 => 'avm,token_table',
		195 => 'avm,token_index',
		211 => 'avm,num_syms',
		224 => 'avm,addresses',
		238 => 'avm,relative_base',
		256 => 'avm,offsets',
		268 => 'avm,data-align',
	},
	# HWR 253 - 6000
	253 => {
		131 => 'avm,kernel_text_start',
		153 => 'avm,names',
		163 => 'avm,token_table',
		179 => 'avm,token_index',
		195 => 'avm,num_syms',
		208 => 'avm,addresses',
		222 => 'avm,relative_base',
		240 => 'avm,offsets',
		252 => 'avm,data-align',
	},
);

# per model hunk config
my %hunkcfg = (
	# HWR 257 - 5530
	257 => {
	},
	# HWR 256 - 7530ax
	256 => {
		'kernel-args' => 117,
	},
	# HWR 253 - 6000
	253 => {
		'kernel-args' => 267,
	},
);

# catch the terminating funktions
if (defined($ARGV[0])) {
	help() if $ARGV[0] =~ /^(-\?|--help)$/;
	version() if $ARGV[0] =~ /^(-v|--version)$/;
	release() if $ARGV[0] eq '--release'; # create release archive - BUMP $VERSION IN SECOND LINE AFTER!
}

# read the arguments: -l -t -x -r -c -s -o -d require an argument, -n -f -h and -q do not
my %opts;
getopts('l:t:x:r:c:s:o:d:nfhq', \%opts) or abort("Invalid arguments\nSee 'fitimg --help'");
abort("nothing to do - you need to pass one of -l -t -x -r or -s\nSee 'fitimg --help'") if !$opts{l} && !$opts{t} && !$opts{x} && !$opts{r} && !$opts{c} && !$opts{s};
abort("don't mix -l -t -x -r -c and -s - only use one of them\nSee 'fitimg --help'") if defined($opts{l}) + defined($opts{t}) + defined($opts{x}) + defined($opts{r}) + defined($opts{c}) + defined($opts{s}) > 1;
abort("replace (-r) and copy (-c) require an output file (-o)\nSee 'fitimg --help'") if ($opts{r} || $opts{c}) && !$opts{o};
abort("input and output file must not be the same\nSee 'fitimg --help'") if ($opts{r} && ($opts{o} eq $opts{r})) || ($opts{c} && ($opts{o} eq $opts{c}));

# read the source file into inbuf
my ($inbuf, $opt, $infile, $insize);
foreach $opt ('l', 't', 'x', 'r', 'c', 's') {
	$infile = $opts{$opt} if defined($opts{$opt});
}
open(INFILE, $infile) || abort("Can't open $infile");
binmode INFILE;
$insize = -s $infile;
read(INFILE, $inbuf, $insize);
close(INFILE);

# find the signature in inbuf
my $pos;
my $offset;
if ($inbuf =~ /\x0D\x00\xED\xFE.{76}\x00\x00\x00\x38/ms) { # 0x0d00edfe AVM fit signature + 0x00000038 @ offset 0x50
	$offset = $-[0];
	$pos = $offset + 0x80;
	my $fitsize = unpack('V', substr($inbuf, $offset + 4, 4)) + 0x50;
	my $fitsize2 = unpack('N', substr($inbuf, $offset + 0x4C, 4)) + 0x50;
	if ($offset > 0) {
		abort("fit-image size mismatch ($fitsize vs $fitsize), can't recover: $infile") if $fitsize != $fitsize2;
		abort("Truncated fit-image: $infile") if ($insize - $offset) < $fitsize;
		$insize = $fitsize + $offset;
		report(sprintf("found $fitsize bytes fit-image at offset 0x%08x", $offset));
	} else {
		abort("Truncated fit-image: $infile") if $insize < $fitsize;
	}
} else {
	abort("Signature not found - not an AVM fit-image: $infile");
}

# verify the passed source/target directory
my $dir = '';
if ((defined($opts{x}) || defined($opts{r})) && defined($opts{d})) {
	$dir = $opts{d};
	abort("Can't find directory: $dir") if !-e $dir;
	abort("Not a directory: $dir") if !-d $dir;
	$dir =~ s|([^/])$|$1/|g; # add trailing slash if missing
}

# now parse the inbuf
my ($head, $type, $body, $strbuf, $string, $name, $len, $status, $blob, $blobsize, $loadaddr, $outbuf, %replaced, $didend);
my $errors = 0;
my $files = 0;
my $lastpos = $pos;
my $oldpos;
my $section = '';
my $realsect = '';
my $hwr = '';

while ($pos < $insize) {
	$oldpos = $pos;
	$head = unpack('N', substr($inbuf, $pos, 4));

	if ($head == 0) {
		$type = unpack('N', substr($inbuf, $pos + 4, 4));
		report('') if $opts{s};
		if ($type == 0) {
			show($oldpos, 8, $head, 'hunk end');
			$pos += 8;
			$didend = 1;
		} else {
			show($oldpos, 8, $head, '???');
			abort(sprintf("unknown hunk end $type @ 0x%08x", $pos));
		}
		if (defined($opts{r}) || defined($opts{c})) {
			$outbuf .= substr($inbuf, $lastpos, $pos - $lastpos);
			$lastpos = $pos;
		}
		last if $didend;

	} elsif ($head == 1) {
		$string = $strbuf = '';
		$pos += 4;
		do {	$strbuf = substr($inbuf, $pos, 4);
			$string .= $strbuf;
			$pos += 4;
		} until $strbuf =~ "\0";
		$string =~ s/\0//g;
		show($oldpos, $pos - $oldpos, $head, "string = '$string'");
		$realsect = $section = $string if $string !~ /^hash_/;
		$hwr = $1 if $section =~ /_HW0*(\d{3})_/;
		if (defined($opts{r}) || defined($opts{c})) {
			$outbuf .= substr($inbuf, $lastpos, $pos - $lastpos);
			$lastpos = $pos;
		}

	} elsif ($head == 2) {
		$string = $strbuf = '';
		$type = unpack('N', substr($inbuf, $pos+4, 4));
		for my $i (0 .. $type-1) { $pos += 4; }
		$pos += 4;
		$pos += 4 if unpack('N', substr($inbuf, $pos, 4)) == 1; # last before configurations
		$pos += 4 if unpack('N', substr($inbuf, $pos, 4)) == 9; # last before description
		do {	$strbuf = substr($inbuf, $pos, 4);
			$string .= $strbuf;
			$pos += 4;
		} until $strbuf =~ "\0";
		$string =~ s/\0//g;
		show($oldpos, $pos - $oldpos, $head, "string = '$string'");
		$realsect = $section = $string if $string !~ /^hash_/;
		$hwr = $1 if $section =~ /_HW0*(\d{3})_/;
		if ($section eq 'description') {
			report('') if $opts{s};
			$strbuf = substr($inbuf, $pos, $insize - $pos - 8);
			$strbuf =~ s/\0$//g;
			my @desc = split(/\0/, $strbuf);
			foreach $string (@desc) { 
				show($pos, length($string) + 1, '->', "array[] = '$string'");
				$pos += length($string) + 1; 
			}
		}
		if (defined($opts{r}) || defined($opts{c})) {
			$outbuf .= substr($inbuf, $lastpos, $pos - $lastpos)    ;
			$lastpos = $pos;
		}

	} elsif ($head == 3) {
		$body = unpack('N', substr($inbuf, $pos + 4, 4));
		$type = unpack('N', substr($inbuf, $pos + 8, 4));

		if ($type == 65) { # blob
			show($oldpos, $body, "$head #$type", "data = <binary blob>");
			report('...') if $opts{h};
			$section = $fname{$hwr}{$section} if defined($opts{f}) && defined($fname{$hwr}{$section});
			$blobsize = $body;
			$blob = substr($inbuf, $pos + 12, $body) if defined($opts{t}) || defined($opts{x}) || defined($opts{c});
			if (defined($opts{r})) {
				if (open(BLOB, "$dir$section")) { # fail silently
					binmode BLOB;
					$blobsize = -s "$dir$section";
					read(BLOB, $blob, $blobsize);
					close BLOB;
					report(sprintf("Replacing: %08d->%08d $section", $body, $blobsize));
					$outbuf .= pack('N', $head);
					$outbuf .= pack('N', $blobsize);
					$outbuf .= pack('N', $type);
					$outbuf .= $blob;
					$outbuf .= "\0" x (4 - $blobsize % 4) if $blobsize % 4; # padding
					$replaced{$section}++;
					$files++;
				}
			}
			delete $missing{$hwr}{$realsect} if $body + $pos < $insize;

		} elsif ($type == 95) { # loadaddr
			$loadaddr = unpack('N', substr($inbuf, $pos + 12, $body));
			show($oldpos, $body, "$head #$type", sprintf("loadaddr = 0x%08x", $loadaddr));
			report(sprintf("%08x-%08x %8d $section", $loadaddr, $loadaddr + $blobsize, $blobsize)) if defined($opts{l});

		} elsif ($type == 111) { # hashvalue (crc)
			my $crc1 = unpack('N', substr($inbuf, $pos + 12, $body));
			show($oldpos, $body, "$head #$type", sprintf("hashvalue = 0x%08x", $crc1));
			if (defined($opts{t}) || defined($opts{x}) || defined($opts{c})) {
				if (defined($opts{x})) {
					if (!defined($opts{n}) || $section !~ /flat_dt|flatdt/) {
						open(BLOB, ">$dir$section") || abort("Can't write to $dir$section");
						binmode BLOB;
						print BLOB $blob;
						close BLOB;
						$files++;
					}
				} else {
					$files++;
				}
				my $crc2 = crc32($blob);
				$status = $crc1 == $crc2 ? 'OK' : 'BAD';
				if (!defined($opts{n}) || $section !~ /flat_dt|flatdt/) {
					report(sprintf("$status: %08x %08x %8d $section", $crc1, $crc2, $blobsize));
				}
				$errors++ if $crc1 != $crc2;
			} elsif (defined($opts{r}) && defined($replaced{$section})) {
					$outbuf .= pack('N', $head);
					$outbuf .= pack('N', 4);
					$outbuf .= pack('N', $type);
					$outbuf .= pack('N', crc32($blob));
			}

		} elsif (($name = $strname{0}{$type}) || ($name = $strname{$hwr}{$type})) { # string
			$string = substr($inbuf, $pos + 12, $body);
			$string =~ s/\0//g;
			show($oldpos, $body, "$head #$type", "$name = '$string'");

			if (defined($opts{r}) && defined($replaced{$section}) && defined($hunkcfg{$hwr}{'kernel-args'}) && ($type == $hunkcfg{$hwr}{'kernel-args'})) { # kernel-args
				my $newstring;
				if (($string =~ /^(.*)0x([0-9a-fA-F]{8}),0x([0-9a-fA-F]{8})(.*:)(\d+)(\@.*)$/)) {
					my $endaddr = hex($2) + $blobsize + ($blobsize % 0x100000 ? (0x100000 - $blobsize % 0x100000) : 0); # 1MB padding
					$newstring = sprintf($1."0x$2,0x%08x$4%d$6", $endaddr, $blobsize); 
				} else {
					$newstring = $string;
				}
				my $newlen = length($newstring) + 1;
				$outbuf .= pack('N', $head);
				$outbuf .= pack('N', $newlen);
				$outbuf .= pack('N', $type);
				$outbuf .= $newstring . "\0";
				$outbuf .= "\0" x (4 - $newlen % 4) if $newlen % 4; # padding
			}

		} elsif (($name = $wordname{0}{$type}) || ($name = $wordname{$hwr}{$type})) { # word
			my $word = unpack('N', substr($inbuf, $pos + 12, $body));
			show($oldpos, $body, "$head #$type", sprintf("$name = 0x%08x", defined($word) ? $word : 0xBADC0DE));

		} else {
			show($oldpos, $body, "$head #$type", 'todo');
		}

		$len = $body + 12 + ($body % 4 ? (4 - $body % 4) : 0); # add hunk head + padding
		$pos += $len;
		if (defined($opts{r}) || defined($opts{c})) {
			$outbuf .= substr($inbuf, $lastpos, $pos - $lastpos) unless defined($replaced{$section}) && ($type == 65 || $type == 111 || (defined($hunkcfg{$hwr}{'kernel-args'}) && ($type == $hunkcfg{$hwr}{'kernel-args'})));
			$lastpos = $pos;
		}

	} else {
		abort(sprintf("unknown hunk head $head @ %08x", $pos));
	}

}

# write the header and outbuf
if (defined($opts{r}) || defined($opts{c})) {
	my $lendiff = length($outbuf) + 0x80 - ($pos - $offset);
	open(OUTFILE, ">$opts{o}") || abort("Can't write to $opts{o}");
	binmode OUTFILE;
	print OUTFILE "\x0D\x00\xED\xFE"; # signature
	print OUTFILE pack('V', $pos - 0x50 - $offset + $lendiff);
	print OUTFILE substr($inbuf, 0x08 + $offset, 0x4C - 0x08);
	print OUTFILE pack('N', $pos - 0x50 - $offset + $lendiff);
	print OUTFILE substr($inbuf, 0x50 + $offset, 4);
	print OUTFILE pack('N', unpack('N', substr($inbuf, 0x54 + $offset, 4)) + $lendiff);
	print OUTFILE substr($inbuf, 0x58 + $offset, 0x6C - 0x58);
	print OUTFILE pack('N', unpack('N', substr($inbuf, 0x6C + $offset, 4)) + $lendiff);
	print OUTFILE substr($inbuf, 0x70 + $offset, 0x80 - 0x70);
	print OUTFILE $outbuf;
	close OUTFILE;
}

# check for truncated files or images with wrong EOF pointers
if ($hwr) {
	my $errs = 0;
	$errs++ unless $didend;
	foreach my $item (sort keys %{$missing{$hwr}}) { 
		print STDERR "### fitimg ### - missing blob '$item'\n";
		$errs++; 
	}
	abort("truncated fit image" . ($offset ? ' or wrong size info' : '')) if $errs;
}

# status reports
if (defined($opts{t}) || defined($opts{x}) || defined($opts{r}) || defined($opts{c})) {
	abort("BROKEN fit image") if $errors;
	if (!defined($opts{q})) {
		print "no errors in $files files\n" if defined($opts{t});
		print "extracted $files files\n" if defined($opts{x});
		print "replaced $files files\n" if defined($opts{r});
		print "no errors copying fit image containing $files files\n" if defined($opts{c});
	}
}

######## subroutines ########

# version or help options
sub version {
	print "fitimg version $VERSION - (C) 2021 Ralf Steines aka Hippie2000 - <metamonk\@yahoo.com>\n";
	print "Handle and manipulate firmware images in AVM /var/tmp/fit-image format. GPLv2+.\n";
	print "Docs and latest version can be found at https://boxmatrix.info/wiki/FIT-Image\n";
	exit unless $_[0];
}

# help option
sub help {
	version(1);
	print <<ENDHELP;

Usage:
  fitimg -l <infile> [-f] [-q]
    List all binaries contained in fit-image <infile>.
    Option -q could be used to silently test the image structure.

  fitimg -t <infile> [-f] [-q]
    Test the integrity of all binaries contained in fit-image <infile>. Performs CRC32 validation.
    Option -q could be used to silently test the image structure and checksum integrity.

  fitimg -x <infile> [-d <dir>] [-n] [-f] [-q]
    Extract all contents of fit-image <infile> or just <file> to current directory or <dir>.
    Option -n suppresses extracting device tree files.
    Option -q suppresses listing which files were extracted.

  fitimg -r <infile> -o <outfile> [-d <dir>] [-f] [-q]
    Replace all contens of fit-image <infile> which exist in current directory or <dir> and write it to <outfile>.
    Files which do not exist in current directory or <dir> will not be replaced.
    Option -q suppresses listing which files were replaced.

  fitimg -c <infile> -o <outfile> [-f] [-q] (fitimg 0.2+)
    Copy an unaltered fit-image from <infile> to <outfile> while testing its integrity.
    This is mainly useful to extract and validate a fit-image from a recovery.exe or firmware.image.
    Option -q could be used to silently copy and test the image structure and checksum integrity.

  fitimg -s <infile> [-h] [-q] (fitimg 0.2+)
    Show the complete hunk structure of the fit-image <infile>.
    Option -h adds a hexdump of all hunks, binaries clipped to 64 bytes
    Option -q could be used to silently test the image structure

Options:
  <infile> can be a fit-image, a firmware.image or a recovery.exe.  (fitimg 0.2+)
  -f activates Freetz mode using filesystem[2].image and kernel[2].image etc instead of the fit names.

  -? (fitimg 0.2+) or --help print this help text and terminates.
  -v (fitimg 0.2+) or --version print this program's version and terminates.

Result:
	Returns 1 on error, otherwise 0.
ENDHELP
	exit;
}

# quietable print
sub report {
	return if defined($opts{q});
	print "$_[0]\n";
}

# error handling
sub abort {
	print STDERR "### fitimg ### - $_[0]\n";
	exit 1;
}

# create release archive
sub release {
	my $release = "fitimg-$VERSION";
	print "Creating release archive $release.tar.gz\n";
	system "mkdir -p ~/hosted/hippie2000/$release";
	system "cp -p ~/bin/fitimg ~/bin/COPYING ~/hosted/hippie2000/$release";
	system "cd ~/hosted/hippie2000 && tar cfvz $release.tar.gz $release";
	if (defined($ARGV[1]) && $ARGV[1] eq '--public') {
		system "chmod 644 ~/hosted/hippie2000/$release.tar.gz";
		system "ln -s -f $release.tar.gz ~/hosted/hippie2000/fitimg-latest.tar.gz";
	} else {
		system "chmod 600 ~/hosted/hippie2000/$release.tar.gz";
	}
	exit;
}

# show hunk - optionally with hexdump (todo)
sub show {
	return if !$opts{s} || $opts{q};
	my ($pos, $len, $hunk, $funct) = @_;
	my $info = "  hunk $hunk [$len] - $funct\n";
	my $count = 0;
	my $str = '';
	print "\n" if $funct =~ /string/ && $funct !~ /hash_/;
	$len = 64 if $len > 64;
	my $pre = ($pos - $offset) % 4;

	if (1) { # FIXME: !$opts{h}
		printf "%08x$info", $pos;

	} else { # FIXME: hexdump still is broken
		print sprintf("%08x ", $pos - $pos % 4);
		if ($pos %4) {
			for my $i (1 .. $pos % 4) { $str .= '..'; }
		}
		while ($count < $len)  {
			$str .= sprintf("%02x", unpack('C', substr($inbuf, $pos + $count, 1)));
			$count++;
			if (($pos + $count) % 4 == 0 && $count <= $len) {
				$str .= $info;
				$str .= sprintf("%08x ", $pos + $count) unless $count == $len;
				$info = "\n";
			}
		}
		if ($pos %4) {
			for my $i (1 .. 4 - ($pos + $len) % 4) { $str .= '..'; }
		}
		$str .= "\n" if $str !~ /\n$/;
		print $str;
	}
}


