#!/usr/bin/env perl

use strict;
use integer;

my $trailer_magic = 'RPTL';

my $tmpfile1 = "/tmp/mkknlimg_$$.1";
my $tmpfile2 = "/tmp/mkknlimg_$$.2";

my $dtok = 0;

while ($ARGV[0] =~ /^-/)
{
    my $arg = shift(@ARGV);
    if ($arg eq '--dtok')
    {
	$dtok = 1;
    }
    else
    {
	print ("* Unknown option '$arg'\n");
	usage();
    }
}

usage() if (@ARGV != 2);

my $kernel_file = $ARGV[0];
my $out_file = $ARGV[1];

if (! -r $kernel_file)
{
    print ("* File '$kernel_file' not found\n");
    usage();
}

my @wanted_config_lines =
(
	'CONFIG_BCM2708_DT'
);

my @wanted_strings =
(
	'bcm2708_fb',
	'brcm,bcm2708-pinctrl',
	'brcm,bcm2835-gpio',
	'of_find_property'
);

my $res = try_extract($kernel_file, $tmpfile1);
$res = try_decompress('\037\213\010', 'xy',    'gunzip', 0,
		      $kernel_file, $tmpfile1, $tmpfile2) if (!$res);
$res = try_decompress('\3757zXZ\000', 'abcde', 'unxz --single-stream', -1,
		      $kernel_file, $tmpfile1, $tmpfile2) if (!$res);
$res = try_decompress('BZh',          'xy',    'bunzip2', 0,
		      $kernel_file, $tmpfile1, $tmpfile2) if (!$res);
$res = try_decompress('\135\0\0\0',   'xxx',   'unlzma', 0,
		      $kernel_file, $tmpfile1, $tmpfile2) if (!$res);
$res = try_decompress('\211\114\132', 'xy',    'lzop -d', 0,
		      $kernel_file, $tmpfile1, $tmpfile2) if (!$res);
$res = try_decompress('\002\041\114\030', 'xy',    'lz4 -d', 1,
		      $kernel_file, $tmpfile1, $tmpfile2) if (!$res);

my $append_trailer;
my $trailer;

$append_trailer = $dtok;

if ($res)
{
    print("Version: $res->{''}\n");

    $append_trailer = $dtok;
    if (!$dtok)
    {
	if (config_bool($res, 'bcm2708_fb'))
	{
	    $dtok ||= config_bool($res, 'CONFIG_BCM2708_DT');
	    $dtok ||= config_bool($res, 'brcm,bcm2708-pinctrl');
	    $dtok ||= config_bool($res, 'brcm,bcm2835-gpio');
	    $append_trailer = 1;
	}
	else
	{
	    print ("* This doesn't look like a Raspberry Pi kernel. In pass-through mode.\n");
	}
    }
}
elsif (!$dtok)
{
    print ("* Is this a valid kernel? In pass-through mode.\n");
}

if ($append_trailer)
{
    printf("DT: %s\n", $dtok ? "y" : "n");

    my @atoms;

    push @atoms, [ $trailer_magic, pack('V', 0) ];
    push @atoms, [ 'KVer', $res->{''} ];
    push @atoms, [ 'DTOK', pack('V', $dtok) ];

    $trailer = pack_trailer(\@atoms);
    $atoms[0]->[1] = pack('V', length($trailer));

    $trailer = pack_trailer(\@atoms);
}

my $ofh;
my $total_len = 0;

if ($out_file eq $kernel_file)
{
    die "* Failed to open '$out_file' for append\n"
	if (!open($ofh, '>>', $out_file));
    $total_len = tell($ofh);
}
else
{
    die "* Failed to open '$kernel_file'\n"
	if (!open(my $ifh, '<', $kernel_file));
    die "* Failed to create '$out_file'\n"
	if (!open($ofh, '>', $out_file));

    my $copybuf;
    while (1)
    {
	my $bytes = sysread($ifh, $copybuf, 64*1024);
	last if (!$bytes);
	syswrite($ofh, $copybuf, $bytes);
	$total_len += $bytes;
    }
    close($ifh);
}

if ($trailer)
{
    # Pad to word-alignment
    syswrite($ofh, "\x000\x000\x000", (-$total_len & 0x3));
    syswrite($ofh, $trailer);
}

close($ofh);

exit($trailer ? 0 : 1);

END {
	unlink($tmpfile1) if ($tmpfile1);
	unlink($tmpfile2) if ($tmpfile2);
}


sub usage
{
	print ("Usage: mkknlimg [--dtok] <vmlinux|zImage|bzImage> <outfile>\n");
	exit(1);
}

sub try_extract
{
	my ($knl, $tmp) = @_;

	my $ver = `strings "$knl" | grep -a -E "^Linux version [1-9]"`;

	return undef if (!$ver);

	chomp($ver);

	my $res = { ''=>$ver };
	my $string_pattern = '^('.join('|', @wanted_strings).')$';

	my @matches = `strings \"$knl\" | grep -E \"$string_pattern\"`;
	foreach my $match (@matches)
	{
	    chomp($match);
	    $res->{$match} = 1;
	}

	my $config_pattern = '^('.join('|', @wanted_config_lines).')=(.*)$';
	my $cf1 = 'IKCFG_ST\037\213\010';
	my $cf2 = '0123456789';

	my $pos = `tr "$cf1\n$cf2" "\n$cf2=" < "$knl" | grep -abo "^$cf2"`;
	if ($pos)
	{
		$pos =~ s/:.*[\r\n]*$//s;
		$pos += 8;
		my $err = (system("tail -c+$pos \"$knl\" | zcat > $tmp 2> /dev/null") >> 8);
		if (($err == 0) || ($err == 2))
		{
			if (open(my $fh, '<', $tmp))
			{
				while (my $line = <$fh>)
				{
					chomp($line);
					$res->{$1} = $2 if ($line =~ /$config_pattern/);
				}

				close($fh);
			}
		}
	}

	return $res;
}


sub try_decompress
{
	my ($magic, $subst, $zcat, $idx, $knl, $tmp1, $tmp2) = @_;

	my $pos = `tr "$magic\n$subst" "\n$subst=" < "$knl" | grep -abo "^$subst"`;
	if ($pos)
	{
		chomp($pos);
		$pos = (split(/[\r\n]+/, $pos))[$idx];
		$pos =~ s/:.*[\r\n]*$//s;
		my $cmd = "tail -c+$pos \"$knl\" | $zcat > $tmp2 2> /dev/null";
		my $err = (system($cmd) >> 8);
		return undef if (($err != 0) && ($err != 2));

		return try_extract($tmp2, $tmp1);
	}

	return undef;
}

sub pack_trailer
{
	my ($atoms) = @_;
	my $trailer = pack('VV', 0, 0);
	for (my $i = $#$atoms; $i>=0; $i--)
	{
		my $atom = $atoms->[$i];
		$trailer .= pack('a*x!4Va4', $atom->[1], length($atom->[1]), $atom->[0]);
	}
	return $trailer;
}

sub config_bool
{
	my ($configs, $wanted) = @_;
	return (($configs->{$wanted} eq 'y') || ($configs->{$wanted} eq '1'));
}