#!/usr/bin/perl
#
# Set PXELINUX hard-coded options
#

use Socket;			# For gethostbyname
use Fcntl;
use bytes;

%option_names = (
      6 => 'domain-name-servers',
     15 => 'domain-name',
     54 => 'next-server',
    209 => 'config-file',
    210 => 'path-prefix',
    211 => 'reboottime'
    );

@fmt_oneip   = ("ip-address", \&parse_oneip, \&show_ip);
@fmt_multiip = ("ip-address-list", \&parse_multiip, \&show_ip);
@fmt_string  = ("string", \&parse_string, \&show_string);
@fmt_uint32  = ("uint32", \&parse_uint32, \&show_uint32);

%option_format = (
      6 => \@fmt_multiip,
     15 => \@fmt_string,
     54 => \@fmt_oneip,
     67 => \@fmt_string,
    209 => \@fmt_string,
    210 => \@fmt_string,
    211 => \@fmt_uint32
    );

sub parse_oneip($)
{
    my($s) = @_;
    my($name,$aliases,$addrtype,$length,@addrs) = gethostbyname($s);

    return ($addrtype == AF_INET) ? $addrs[0] : undef;
}

sub parse_multiip($)
{
    my($l) = @_;
    my $s;
    my @a = ();
    my $addr;
    my $d = '';

    foreach $s (split(/,/, $l)) {
	my($name,$aliases,$addrtype,$length,@addrs)
	    = gethostbyname($s);
	if ($addrtype == AF_INET) {
	    foreach $addr (@addrs) {
		$d .= $addr;
	    }
	}
    }

    return $d ne '' ? $d : undef;
}

sub show_ip($)
{
    my($l) = @_;

    if (length($l) & 3) {
	return undef;
    } else {
	my @h = ();
	my $i;

	for ($i = 0; $i < length($l); $i += 4) {
	    push(@h, inet_ntoa(substr($l, $i, 4)));
	}

	return join(',', @h);
    }
}

sub parse_string($)
{
    return $_[0];
}

sub show_string($)
{
    my($s) = @_;
    my $o, $i, $c;

    $o = "\'";
    for ($i = 0; $i < length($s); $i++) {
	$c = substr($s, $i, 1);
	if ($c eq "\'" || $c eq '!') {
	    $o .= "\'\\$c\'";
	} else {
	    $o .= $c;
	}
    }
    $o .= "\'";

    return $o;
}

sub parse_uint32($)
{
    my($s) = @_;

    if ($s =~ /^[0-9]+$/) {
	return pack("N", $s);
    } else {
	return undef;
    }
}

sub show_uint32($)
{
    my($l) = @_;

    if (length($l) == 4) {
	return unpack("N", $l);
    } else {
	return undef;
    }
}

sub parse_generic($)
{
    my($s) = @_;

    if ($s =~ /^[0-9a-f]{1,2}(:[0-9a-f]{1,2})*$/) {
	my $h;
	my @b = ();

	foreach $h (split(/\:/, $s)) {
	    push(@b, hex $h);
	}

	return pack("C", @b);
    } else {
	return undef;
    }
}

sub show_generic($)
{
    my($l) = @_;
    my $i;
    my @h;

    for ($i = 0; $i < length($l); $i++) {
	push(@h, sprintf("%02x", unpack("C", substr($l, $i, $1))));
    }

    return join(':', @h);
}

sub parse_option($$)
{
    my($opt, $arg) = @_;
    my $v;

    if (defined($option_format{$opt})) {
	$v = $option_format{$opt}[1]($arg);
	return $v if (defined($v));
    }

    return parse_generic($arg);
}

sub show_option($$)
{
    my($opt, $arg) = @_;
    my $v;

    if (defined($option_format{$opt})) {
	$v = $option_format{$opt}[2]($arg);
	return $v if (defined($v));
    }

    return show_generic($arg);
}

sub option_number($)
{
    my($n) = @_;

    if (defined($option_rnames{$n})) {
	return $option_rnames{$n};
    } elsif ($n =~ /^[0-9]+$/ && $n >= 1 && $n <= 254) {
	return $n+0;
    } else {
	return undef;
    }
}

sub read_optsets($)
{
    my($file) = @_;
    my $data, $bdata, $adata;
    my $patch_start = (stat($file))[7];
    my $hdroffset = 0;		# 0 means non-deep-embedded
    my $bufsize = 0;
    my $junk;
    my %hdr;

    return undef unless (seek($file, 0, SEEK_SET));
    return undef unless (read($file, $data, 48) == 48);

    my($mzmagic, $junk, $magic, $len, $flags, $boff, $blen, $aoff, $alen)
	= unpack("va[6]VVVVVVV", $data);

    if ($mzmagic == 0x5a4d) {
	# It is an EFI file... search for the magic number
	$hdroffset = 48;
	my $magic = pack("VVVV", 0x2a171ead, 0x0600e65e,
			 0x4025a4e4, 0x42388fc8);

	while (1) {
	    return undef unless (read($file, $data, 16) == 16);
	    last if ($data eq $magic);

	    $hdroffset += 16;
	}

	return undef unless (read($file, $data, 16) == 16);
	($blen, $alen, $bufsize, $junk) = unpack("VVVV", $data);

	$patch_start = $boff = $hdroffset + 32;
	$aoff = $boff + $blen;

	$hdr{'deep'} = 1;
	$hdr{'bufsize'} = $bufsize;
	$hdr{'hdroffset'} = $hdroffset;
    } else {
	# It is a BIOS PXE file

	return undef if ($magic != 0x2983c8ac);
	return undef if ($len < 7*4);

	$hdr{'deep'} = 0;
    }

    if ($blen == 0) {
	$bdata = '';
    } else {
	return undef unless (seek($file, $boff, SEEK_SET));
	return undef unless (read($file, $bdata, $blen) == $blen);
	$patch_start = $boff if ($boff < patch_start);
    }

    if ($alen == 0) {
	$adata = '';
    } else {
	return undef unless (seek($file, $aoff, SEEK_SET));
	return undef unless (read($file, $adata, $alen) == $alen);
	$patch_start = $aoff if ($aoff < $patch_start);
    }

    $hdr{'patch_start'} = $patch_start;

    return (\%hdr, $bdata, $adata);
}

sub write_optsets($$@)
{
    my($file, $hdr, $bdata, $adata) = @_;
    my $boff = 0;
    my $aoff = 0;
    my $bufsize = 0;
    my $patch_start = $hdr->{'patch_start'};
    my $len;

    $bdata .= "\xff" unless ($bdata eq '');
    $adata .= "\xff" unless ($adata eq '');

    $len = length($bdata) + length($adata);

    if (defined($hdr->{'bufsize'})) {
	return undef unless ($len <= $hdr->{'bufsize'});
    }

    return undef unless (seek($file, $patch_start, SEEK_SET));

    if (length($bdata)) {
	$boff = $patch_start;
	return undef unless (print $file $bdata);
	$patch_start += length($bdata);
    }

    if (length($adata)) {
	$aoff = $patch_start;
	return undef unless (print $file $adata);
	$patch_start += length($adata);
    }

    if ($hdr->{'deep'}) {
	return undef unless (print $file "\0" x ($hdr->{'bufsize'} - $len));
	return undef unless (seek($file, $hdr->{'hdroffset'} + 16, SEEK_SET));
	my $hdr = pack("VV", length($bdata), length($adata));
	return undef unless (print $file $hdr);
    } else {
	my $hdr = pack("VVVV", $boff, length($bdata), $aoff, length($adata));

	return undef unless (seek($file, 8+3*4, SEEK_SET));
	return undef unless (print $file $hdr);

	truncate($file, $patch_start);
    }

    return 1;
}

sub delete_option($$)
{
    my ($num, $block) = @_;
    my $o, $l, $c, $x;

    $x = 0;
    while ($x < length($block)) {
	($o, $l) = unpack("CC", substr($block, $x, 2));
	if ($o == $num) {
	    # Delete this option
	    substr($block, $x, $l+2) = '';
	} elsif ($o == 0) {
	    # Delete a null option
	    substr($block, $x, 1) = '';
	} elsif ($o == 255) {
	    # End marker - truncate block
	    $block = substr($block, 0, $x);
	    last;
	} else {
	    # Skip to the next option
	    $x += $l+2;
	}
    }

    return $block;
}

sub add_option($$$)
{
    my ($num, $data, $block) = @_;

    $block = delete_option($num, $block);

    if (length($data) == 0) {
	return $block;
    } elsif (length($data) > 255) {
	die "$0: option $num has too much data (max 255 bytes)\n";
    } else {
	return $block . pack("CC", $num, length($data)) . $data;
    }
}

sub list_options($$)
{
    my($pfx, $data) = @_;
    my $x, $o, $l;

    while ($x < length($data)) {
	($o, $l) = unpack("CC", substr($data, $x, 2));

	if ($o == 0) {
	    $x++;
	} elsif ($o == 255) {
	    last;
	} else {
	    my $odata = substr($data, $x+2, $l);
	    last if (length($odata) != $l); # Incomplete option

	    printf "%s%-20s %s\n", $pfx,
		$option_names{$o} || sprintf("%d", $o),
		show_option($o, $odata);

	    $x += $l+2;
	}
    }
}

sub usage()
{
    my $i;

    print STDERR "Usage: $0 options pxelinux.0\n";
    print STDERR "Options:\n";
    print STDERR "--before option value   -b   Add an option before DHCP data\n";
    print STDERR "--after  option value   -a   Add an option after DHCP data\n";
    print STDERR "--delete option         -d   Delete an option\n";
    print STDERR "--list                  -l   List set options\n";
    print STDERR "--dry-run               -n   Don't modify the target file\n";
    print STDERR "--help                  -h   Display this help text\n";
    print STDERR "\n";
    print STDERR "The following DHCP options are currently recognized:\n";
    printf STDERR "%-23s %-3s  %s\n", 'Name', 'Num', 'Value Format';

    foreach $i (sort { $a <=> $b } keys(%option_names)) {
	printf STDERR "%-23s %3d  %s\n",
		$option_names{$i}, $i, $option_format{$i}[0];
    }
}

%option_rnames = ();
foreach $opt (keys(%option_names)) {
    $option_rnames{$option_names{$opt}} = $opt;
}

%before   = ();
%after    = ();
@clear    = ();
$usage    = 0;
$err      = 0;
$list     = 0;
$no_write = 0;
undef $file;

while (defined($opt = shift(@ARGV))) {
    if ($opt !~ /^-/) {
	if (defined($file)) {
	    $err = $usage = 1;
	    last;
	}
	$file = $opt;
    } elsif ($opt eq '-b' || $opt eq '--before') {
	$oname = shift(@ARGV);
	$odata = shift(@ARGV);

	if (!defined($odata)) {
	    $err = $usage = 1;
	    last;
	}

	$onum = option_number($oname);
	if (!defined($onum)) {
	    print STDERR "$0: unknown option name: $oname\n";
	    $err = 1;
	    next;
	}

	$odata = parse_option($onum, $odata);
	if (!defined($odata)) {
	    print STDERR "$0: unable to parse data for option $oname\n";
	    $err = 1;
	    next;
	}

	delete $after{$onum};
	$before{$onum} = $odata;
	push(@clear, $onum);
    } elsif ($opt eq '-a' || $opt eq '--after') {
	$oname = shift(@ARGV);
	$odata = shift(@ARGV);

	if (!defined($odata)) {
	    $err = $usage = 1;
	    last;
	}

	$onum = option_number($oname);
	if (!defined($onum)) {
	    print STDERR "$0: unknown option name: $oname\n";
	    $err = 1;
	    next;
	}

	$odata = parse_option($onum, $odata);
	if (!defined($odata)) {
	    print STDERR "$0: unable to parse data for option $oname\n";
	    $err = 1;
	    next;
	}

	delete $before{$onum};
	$after{$onum} = $odata;
	push(@clear, $onum);
    } elsif ($opt eq '-d' || $opt eq '--delete') {
	$oname = shift(@ARGV);

	if (!defined($oname)) {
	    $err = $usage = 1;
	    last;
	}

	$onum = option_number($oname);
	if (!defined($onum)) {
	    print STDERR "$0: unknown option name: $oname\n";
	    $err = 1;
	    next;
	}

	push(@clear, $onum);
	delete $before{$onum};
	delete $after{$onum};
    } elsif ($opt eq '-n' || $opt eq '--no-write' || $opt eq '--dry-run') {
	$no_write = 1;
    } elsif ($opt eq '-l' || $opt eq '--list') {
	$list = 1;
    } elsif ($opt eq '-h' || $opt eq '--help') {
	$usage = 1;
    } else {
	print STDERR "Invalid option: $opt\n";
	$err = $usage = 1;
    }
}

if (!defined($file) && !$usage) {
    $err = $usage = 1;
}
if ($usage) {
    usage();
}
if ($err || $usage) {
    exit($err);
}

if (!scalar(@clear)) {
    $no_write = 1;		# No modifications requested
}

$mode = $no_write ? '<' : '+<';

open(FILE, $mode, $file)
    or die "$0: cannot open: $file: $!\n";
($hdrinfo, @data) = read_optsets(\*FILE);
if (!defined($hdrinfo)) {
    die "$0: $file: patch block not found or file corrupt\n";
}

foreach $o (@clear) {
    $data[0] = delete_option($o, $data[0]);
    $data[1] = delete_option($o, $data[1]);
}
foreach $o (keys(%before)) {
    $data[0] = add_option($o, $before{$o}, $data[0]);
}
foreach $o (keys(%after)) {
    $data[1] = add_option($o, $after{$o}, $data[1]);
}

if ($list) {
    list_options('-b ', $data[0]);
    list_options('-a ', $data[1]);
}

if (!$no_write) {
    if (!write_optsets(\*FILE, $hdrinfo, @data)) {
	die "$0: $file: failed to write options: $!\n";
    }
}

close(FILE);
exit 0;