#! /usr/bin/perl

# Add/remove key=value pairs from iso application area (0x200 bytes starting
# at 0x8373). Entries are separated by semicolons (';').
#
# Digest is calculated assuming all zeros in 0x0000-0x01ff (MBR) and all
# spaces in 0x8373-0x8572.

use strict;

use Getopt::Long;
use Digest::MD5;
use Digest::SHA;

use Data::Dumper;
$Data::Dumper::Sortkeys = 1;
$Data::Dumper::Terse = 1;
$Data::Dumper::Indent = 1;

# Here are some ISO9660 file system related constants.
#
# If in doubt about the ISO9660 layout, check http://alumnus.caltech.edu/~pje/iso9660.html
#
use constant {
  # offset of volume descriptor
  # usually points to "\x01CD001\x01\x00"
  ISO9660_MAGIC_START => 0x8000,

  # offset of volume size (in 2 kB units)
  # stored as 32 bit little-endian, followed by a 32 bit big-endian value
  ISO9660_VOLUME_SIZE => 0x8050,

  # application identifier length
  ISO9660_APP_ID_LENGTH => 0x80,

  # offset of application specific data (anything goes)
  ISO9660_APP_DATA_START => 0x8373,

  # application specific data length
  ISO9660_APP_DATA_LENGTH => 0x200,
};

# Here are some MBR (partition table) related constants.
#
# If in doubt about the MBR layout, check https://wiki.osdev.org/Partition_Table
#
use constant {
  # MBR start
  MBR_OFFSET => 0,

  # offset of MBR magic
  # points to "\x55\xaa"
  MBR_MAGIC_START => 0x1fe,

  # offset of partition table in MBR
  MBR_PARTITION_TABLE => 0x1be,

  # MBR size (in bytes)
  MBR_LENGTH => 0x200,
};

# Some SUSE specific constants.
#
# see
#   https://github.com/openSUSE/mksusecd/blob/master/mksusecd
#   https://github.com/OSInside/kiwi/blob/master/kiwi/iso_tools/iso.py
#
use constant {
  # signature block starts with this string
  SIGNATURE_MAGIC => "7984fc91-a43f-4e45-bf27-6d3aa08b24cf"
};

sub usage;
sub read_image_blob;
sub get_image_type;
sub get_padding;
sub read_tags;
sub write_tags;
sub parse_tag;
sub get_tag;
sub set_tag;
sub remove_tag;
sub prepare_buffer;
sub add_to_digest;
sub calculate_digest;
sub export_tags;
sub export_signature;
sub import_signature;
sub normalize_buffer;

my $opt_digest = undef;
my $opt_check = 0;
my $opt_pad = undef;
my $opt_show = 1;
my $opt_clean = 0;
my @opt_add_tag;
my @opt_remove_tag;
my $opt_verbose;
my $opt_tags_export;
my $opt_signature_export;
my $opt_signature_import;

GetOptions(
  'show'               => \$opt_show,
  'md5|md5sum'         => sub { $opt_digest = 'md5' },
  'digest=s'           => \$opt_digest,
  'check'              => \$opt_check,
  'pad=i'              => \$opt_pad,
  'add-tag=s'          => \@opt_add_tag,
  'remove-tag=s'       => \@opt_remove_tag,
  'export-tags=s'      => \$opt_tags_export,
  'export-signature=s' => \$opt_signature_export,
  'import-signature=s' => \$opt_signature_import,
  'clean'              => \$opt_clean,
  'verbose|v'          => \$opt_verbose,
  'help'               => sub { usage 0 },
);

my $image_data;			# hash ref with image related data
my $current_tags = [];		# current list of tags
my $old_tags = [];		# original list of tags
my $digest_iso;			# digest calculated over iso image
my $digest_part;		# digest calculated over partition

# Note: all '*_block' variables use 0.5 kB units.


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
$image_data->{name} = shift;
if($image_data->{name} eq '') {
  print STDERR "tagmedia: no image specified\n";
  usage 1;
}

$image_data->{write} = $opt_digest || defined $opt_pad || $opt_check || @opt_add_tag || @opt_remove_tag || $opt_clean;

if($opt_digest =~ /^md5(sum)?$/i) {
  $digest_iso = Digest::MD5->new;
  $digest_part = Digest::MD5->new;
}
elsif($opt_digest =~ /^sha(1|224|256|384|512)(sum)?$/i) {
  $digest_iso = Digest::SHA->new($1);
  $digest_part = Digest::SHA->new($1);
}
elsif($opt_digest) {
  die "$opt_digest: unsupported digest\n";
}

read_image_blob $image_data;
get_image_type $image_data;

$old_tags = read_tags $image_data unless $opt_clean;

# clone tag list
set_tag $current_tags, { key => $_->{key}, value =>  $_->{value} } for @$old_tags;

set_tag $current_tags, { key => "check", value => 1 } if $opt_check;

get_padding $image_data, $current_tags;

if($opt_tags_export) {
  export_tags $image_data, $opt_tags_export;

  exit 0;
}

if($opt_signature_export) {
  export_signature $image_data, $opt_signature_export;

  exit 0;
}

if($opt_signature_import) {
  import_signature $image_data, $opt_signature_import;

  exit 0;
}

if(my $sig = get_tag $current_tags, "signature") {
  $image_data->{signature_start} = $sig->{value} + 0;
}

prepare_buffer $image_data;

# print Dumper $image_data;
# print Dumper $current_tags;

# calculate digest
calculate_digest $image_data if $opt_digest;

# finally close file handle (had been opened in read_image_blob())
close $image_data->{fh};

if($image_data->{signature_start}) {
  set_tag $current_tags, { key => "signature", value => $image_data->{signature_start} };
}

for (@opt_add_tag) {
  set_tag $current_tags, parse_tag($_);
}

for (@opt_remove_tag) {
  $current_tags = remove_tag $current_tags, $_;
}

write_tags $image_data, $current_tags if $image_data->{write};

# print tags for user's pleasure
for (@{$image_data->{write} ? $current_tags : $old_tags}) {
  print "$_->{key}";
  print " = $_->{value}" if defined $_->{value};
  print "\n";
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# usage(exit_code)
#
# Print help message and exit program with exit_code.
#
sub usage
{
  (my $msg = <<"  = = = = = = = =") =~ s/^ {4}//mg;
    Usage: tagmedia [OPTIONS] IMAGE
    Add/remove tags to SUSE installation media.

    IMAGE is an SUSE installation medium; either DVD image or disk image.

    Options:

      --show                    Show currently set tags (default if no option is given).
      --digest DIGEST           Add digest DIGEST (md5, sha1, sha224, sha256, sha384, sha512).
      --pad N                   Ignore N 2 kB blocks of padding (at image end).
      --check                   Tell installer to run media check at startup.
      --add-tag foo=bar         Add tag foo with value bar.
      --remove-tag foo          Remove tag foo.
      --export-tags FILE        Export raw tag data to FILE.
      --export-signature FILE   Export image signature to FILE.
      --import-signature FILE   Import image signature from FILE.
      --clean                   Remove all tags.
      --help                    Write this help text.

  = = = = = = = =

  print $msg;

  exit shift;
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# read_image_blob(image)
#
# Read first 36 kB.
# Note: leaves $image_data ->{fh} open for calculate_digest() to continue reading.
#
# image: hash with image related data
#
sub read_image_blob
{
  my ($image) = @_;

  my $blob_size = 0x9000;
  $image->{blob_blocks} = $blob_size >> 9;		# in 0.5 kB units

  die "$image->{name}: $!\n" unless open $image->{fh}, $image->{name};
  die "$image->{name}: $!\n" unless sysread $image->{fh}, $image->{blob}, $blob_size;
  die "$image->{name}: file too short\n" if $blob_size != length $image->{blob};
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# get_image_type(image)
#
# Analyze image type and get some basic data.
#
# image: hash with image related data
# image->{blob}: buffer containing a sufficiently large portion of the image (36 kB)
#
sub get_image_type
{
  my ($image) = @_;

  # Get iso size from ISO9660 header.
  #
  # Don't verify too much - most of the header might be gone. Since in
  # ISO9660 sizes are stored both-endian we can do a plausibility check
  # which should be enough.
  #
  my $iso_magic_ok = substr($image->{blob}, ISO9660_MAGIC_START, 8) eq "\x01CD001\x01\x00";
  my $little = 4 * unpack("V", substr($image->{blob}, ISO9660_VOLUME_SIZE, 4));
  my $big = 4 * unpack("N", substr($image->{blob}, ISO9660_VOLUME_SIZE + 4, 4));
  if(
    $iso_magic_ok &&
    $little &&
    $little == $big &&
    $little <= (-s $image->{name}) / 512
  ) {
    $image->{iso_blocks} = $little;
  }

  # Scan mbr for last primary partition.
  #
  # Set $image->{part_start}, $image->{part_blocks} (in 0.5 kB units).
  #
  if(substr($image->{blob}, MBR_MAGIC_START, 2) eq "\x55\xaa") {
    for (my $idx = 0; $idx < 4; $idx++) {
      my ($boot, $type, $start, $size) = unpack(
        "Cx3Cx3VV",
        substr($image->{blob}, MBR_PARTITION_TABLE + 0x10 * $idx, 0x10)
      );
      if(
        $type &&
        !($boot & 0x7f) &&
        $size &&
        $start + $size > $image->{part_start} + $image->{part_blocks}
      ) {
        $image->{part_start} = $start;
        $image->{part_blocks} = $size;
      }
    }
  }

  print "iso blocks = $image->{iso_blocks}, partition blocks = $image->{part_blocks} @ $image->{part_start}\n" if $opt_verbose;

  die "$image->{name}: unsupported image format\n" unless $image->{iso_blocks} || $image->{part_blocks};
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# get_padding(image, tags)
#
# Get padding value to use.
#
# image: hash with image related data
# tags: array ref with tag hashes
#
# Note: this uses the value passed via '--pad' option and stores it as
# 'pad' tag.
#
sub get_padding
{
  my ($image, $tags) = @_;
  my $padding_set;

  if($image->{iso_blocks}) {
    set_tag $tags, { key => "pad", value => $opt_pad } if $opt_pad;
    my $pad_tag = get_tag $tags, "pad";

    $image->{pad_blocks} = $pad_tag->{value} << 2 if $pad_tag;
    $image->{iso_blocks} -= $image->{pad_blocks};

    die "padding too large\n" if $image->{iso_blocks} <= 0;
  }
  else {
    remove_tag $tags, "pad";
  }
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# tags = read_tags(image)
#
# Read existing tags from image fragment.
#
# image: hash with image related data
# image->{blob}: buffer containing a sufficiently large portion of the image (36 kB)
# tags: array ref with tags
# tags are hashes with key/value pairs - value is undef if there's no '='
#
sub read_tags
{
  my ($image) = @_;

  my $buf = substr($image->{blob}, ISO9660_APP_DATA_START, ISO9660_APP_DATA_LENGTH);
  die "$image->{name}: unsupported image format\n" unless $buf  =~ /^[0-9A-Za-z_=,;! \x00]{512}$/;
  $buf =~ s/[\s\x00]*$//;

  my $tags = [];

  for my $line (split /;/, $buf) {
    set_tag $tags, parse_tag($line);
  }

  return $tags;
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# write_tags(image, $tags)
#
# Insert tags into image fragment.
# Note: this does not write anything to disk!
#
# image: hash with image related data
# image->{blob}: buffer containing a sufficiently large portion of the image (36 kB)
# tags: array ref with tags
#
sub write_tags
{
  my ($image, $tags) = @_;
  my $buf;

  for my $tag (@$tags) {
    $buf .= ";" if defined $buf;
    $buf .= $tag->{key};
    $buf .= "=$tag->{value}" if defined $tag->{value};
  }

  die "tags too large: \"$buf\"\n" if length($buf) > ISO9660_APP_DATA_LENGTH;
  $buf .= " " x (ISO9660_APP_DATA_LENGTH - length($buf));

  my $len;
  die "$image->{name}: $!\n" unless open $image->{fh}, "+<", $image->{name};
  die "$image->{name}: $!\n" unless seek $image->{fh}, ISO9660_APP_DATA_START, 0;
  die "$image->{name}: $!\n" unless $len = syswrite $image->{fh}, $buf, ISO9660_APP_DATA_LENGTH;
  die "$image->{name}: failed to update image\n" unless $len == ISO9660_APP_DATA_LENGTH;

  close $image->{fh};
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# tag = parse_tag(string)
#
# Parse line ('key=value' format) and return hash ref with key/value elements.
# value is undef if there's no '='.
#
sub parse_tag
{
  my ($line) = @_;

  $line =~ s/^\s*|\s*$//g;

  if($line =~ /^([A-Z_\d\s]+?)\s*=\s*+(.*)$/i) {
    return { key => $1, value => $2 };
  }
  else {
    return { key => $line };
  }
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# tag = get_tag(tags, key)
#
# Get tag from tag array with specified key.
# Returns undef if no tag with such a key exists.
#
# tags: array ref with tag hashes
# key: string
# tag: tag hash ref ('key', 'value' pair)
#
sub get_tag
{
  my ($tags, $key) = @_;

  for my $tag (@$tags) {
    return $tag if $tag->{key} eq $key;
  }

  return undef;
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# set_tag(tags, tag)
#
# Set tag in tags array.
#
# tags: array ref with tag hashes
# tag: tag hash ref ('key', 'value' pair)
#
sub set_tag
{
  my ($tags, $tag) = @_;

  my $old_tag = get_tag $tags, $tag->{key};
  if($old_tag) {
    $old_tag->{value} = $tag->{value};
  }
  else {
    push @$tags, $tag;
  }
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# tags = remove_tag(tags, key)
#
# Remove tag from tag array with specified key.
# Returns new tag array.
#
# tags: array ref with tag hashes
# key: string
#
sub remove_tag
{
  my ($tags, $key) = @_;

  return [ grep { $_->{key} ne $key } @$tags ];
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# prepare_buffer(image)
#
# Reset some areas in image blob to defined values before starting digest calculation.
#
# image: hash with image related data
# image->{blob}: buffer containing a sufficiently large portion of the image (36 kB)
#
sub prepare_buffer
{
  my $image = $_[0];

  substr($image->{blob}, MBR_OFFSET, MBR_LENGTH) = "\x00" x MBR_LENGTH;
  substr($image->{blob}, ISO9660_APP_DATA_START, ISO9660_APP_DATA_LENGTH) = " " x ISO9660_APP_DATA_LENGTH;
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# process_digest(digest, start, blocks, buf_start, buf_blocks, buf)
#
# Update digest.
#
# The total digest needs to be calculated from start to start + blocks. The
# current buffer in buf holds file data from buf_start to buf_start + buf_blocks.
#
# digest: digest ref
# start: region start block
# blocks: region blocks
# buf_start: buffer start block
# buf_blocks: buffer blocks
# buf: buffer ref
#
# Note: a block is 0.5 kB.
#
sub process_digest
{
  my ($digest, $start, $blocks, $buf_start, $buf_blocks, $buf) = @_;

  my $end = $start + $blocks;
  my $buf_end = $buf_start + $buf_blocks;

  return if $start >= $buf_end || $end <= $buf_start;

  if($start > $buf_start) {
    my $skip = $start - $buf_start;
    $buf_blocks -= $skip;
    $buf_start = $start;
    $buf = substr($buf, $skip << 9);
  }

  if($buf_end > $end) {
    $buf_blocks -= $buf_end - $end;
    $buf_end = $end;
    $buf = substr($buf, 0, $buf_blocks << 9);
  }

  $digest->add($buf);
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# calculate_digest($image)
#
# Calculate digest.
#
# Note: $image_data->{fh} has been opened in read_image_blob(). We just continue
# reading here.
#
# This function calculates 2 different digests:
#
#   (1) $image->{iso_blocks} starting from 0, plus $image->{pad_blocks} of zeros
#   (2) $image->{part_blocks} starting from $image->{part_start}
#
# $image->{blob_blocks} have already been read and are cached in $image->{blob}.
#
# Note: padding has been subtracted from $image->{iso_blocks}.
#
sub calculate_digest
{
  my ($image) = @_;

  my $iso_start = 0;
  my $iso_blocks = $image->{iso_blocks};
  my $part_start = $image->{part_start};
  my $part_blocks = $image->{part_blocks};

  my $full_blocks = $part_start + $part_blocks;
  $full_blocks = $iso_blocks if $iso_blocks > $full_blocks;

  my $pos = $image->{blob_blocks};
  my $step_blocks = 2048;	# 1 MiB chunks

  process_digest $digest_iso, $iso_start, $iso_blocks, 0, $pos, $image->{blob};
  process_digest $digest_part, $part_start, $part_blocks, 0, $pos, $image->{blob};

  while($pos < $full_blocks) {
    my $buf;
    my $to_read = $full_blocks - $pos;
    $to_read = $step_blocks if $step_blocks < $to_read;

    my $read_len = sysread $image->{fh}, $buf, $to_read << 9;
    die "$image->{name}: read error: $to_read blocks @ $pos\n" if $read_len != $to_read << 9;

    normalize_buffer $image, $pos, \$buf;

    process_digest $digest_iso, $iso_start, $iso_blocks, $pos, $to_read, $buf;
    process_digest $digest_part, $part_start, $part_blocks, $pos, $to_read, $buf;

    $pos += $to_read;
  }

  if($image->{iso_blocks}) {
    my $buf = "\x00" x 0x200;
    my $pad_blocks = $image->{pad_blocks};

    while($pad_blocks-- > 0) {
      $digest_iso->add($buf);
    }

    $digest_iso = $digest_iso->hexdigest;
    set_tag $current_tags, { key => "${opt_digest}sum", value => $digest_iso };
  }

  if($image->{part_blocks}) {
    $digest_part = $digest_part->hexdigest;
    set_tag $current_tags, {
      key => "partition",
      value => "$image->{part_start},$image->{part_blocks},$digest_part"
    }
  }
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# export_tags(image, file)
#
# Export raw tag data from image to file.
#
# image: hash with image related data
# file: file name
#
sub export_tags
{
  my ($image, $file) = @_;

  my $buf = substr($image->{blob}, ISO9660_APP_DATA_START, ISO9660_APP_DATA_LENGTH);

  if(open my $f, ">", $file) {
    print $f $buf;
    close $f;
  }
  else {
    die "$file: $!\n";
  }
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# export_signature(image, file)
#
# Export signature data from image to file.
#
# image: hash with image related data
# file: file name
#
sub export_signature
{
  my ($image, $file) = @_;
  my $sig = get_tag $current_tags, "signature";

  my $buf;

  die "$image->{name}: no signature location found\n" if !$sig || $sig->{value} == 0;
  die "$image->{name}: $!\n" unless seek $image->{fh}, $sig->{value} * 0x200, 0;
  die "$image->{name}: $!\n" unless 0x800 == sysread $image->{fh}, $buf, 0x800;

  die "$image->{name}: invalid signature block\n" if SIGNATURE_MAGIC ne substr($buf, 0, length SIGNATURE_MAGIC);

  substr($buf, 0, 0x40) = "";

  $buf =~ s/\x00*$//;

  if(open my $f, ">", $file) {
    print $f $buf;
    close $f;
  }
  else {
    die "$file: $!\n";
  }
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# import_signature(image, file)
#
# Import signature data from file to image.
#
# image: hash with image related data
# file: file name
#
sub import_signature
{
  my ($image, $file) = @_;
  my $sig = get_tag $current_tags, "signature";

  my $buf;

  die "$image->{name}: no signature location found\n" if !$sig || $sig->{value} == 0;
  die "$image->{name}: $!\n" unless seek $image->{fh}, $sig->{value} * 0x200, 0;
  die "$image->{name}: $!\n" unless 0x800 == sysread $image->{fh}, $buf, 0x800;

  die "$image->{name}: invalid signature block\n" if SIGNATURE_MAGIC ne substr($buf, 0, length SIGNATURE_MAGIC);

  my $buf2;

  if(open my $f, "<", $file) {
    local $/;
    $buf2 = <$f>;
    close $f;
  }
  else {
    die "$file: $!\n";
  }

  if(length($buf2) > 0x800 - 0x40) {
    die "$file: signature too large\n";
  }

  $buf = substr($buf, 0, 0x40) . $buf2;
  $buf .= "\x00" x (0x800 - length($buf));

  die "$image->{name}: $!\n" unless open $image->{fh}, "+<", $image->{name};
  die "$image->{name}: $!\n" unless seek $image->{fh}, $sig->{value} * 0x200, 0;
  die "$image->{name}: $!\n" unless 0x800 == syswrite $image->{fh}, $buf, 0x800;
  close $image->{fh};
}


# - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
# normalize_buffer(image, pos, buffer_ref)
#
# If buffer contains signature block (2 kiB), clear signature block.
# A cleared signature block containts 0x40 bytes magic header, the rest is
# all zeros (0).
#
# This function looks for a signature block and returns its position in
# $image->{signature_start} if $image->{signature_start} is unset.
#
# image: hash with image related data
# pos: block number of buffer start
# buffer_ref: reference to buffer; buffer may get modified
#
sub normalize_buffer
{
  my ($image, $pos, $buf_ref) = @_;
  my $blocks = length($$buf_ref) >> 9;

  if(!$image->{signature_start}) {
    for (my $i = 0; $i < $blocks; $i++) {
      if(SIGNATURE_MAGIC eq substr($$buf_ref, $i << 9, length SIGNATURE_MAGIC)) {
        $image->{signature_start} = $pos + $i;
      }
    }
  }

  if($image->{signature_start}) {
    my $signature_len = 4;	# 2 kiB
    my $x = $image->{signature_start} - $pos;
    if($x >= 0 && $x < $blocks) {
      for (my $i = 0; $i < $signature_len && $x + $i < $blocks; $i++) {
        if($i == 0) {
          # leave 0x40 bytes intact in first block
          substr($$buf_ref, ($x << 9) + 0x40, 0x200 - 0x40) = "\x00" x (0x200 - 0x40);
        }
        else {
          substr($$buf_ref, ($x + $i) << 9, 0x200) = "\x00" x 0x200;
        }
      }
    }
  }
}
