#!/usr/bin/perl -c

# Module to handle tar files streams.  This module assumes that neither
# stream is seekable, but can safely parse items on and off a tar stream
# reporting errors.

package ME::TarStream;

use strict;
use warnings;

use IO::File;
use IO::Handle;
use File::Temp;

use constant BLOCKSIZE => 512;

our $VERBOSE = 0;

sub new {
  my $class = shift;
  my %args = @_;
  $args{inpos} ||= 0;
  $args{outpos} ||= 0;
  my $Self = bless {
    %args
  }, ref($class) || $class;
  return $Self;
}

sub InPos {
  my $Self = shift;
  return $Self->{inpos};
}

sub OutPos {
  my $Self = shift;
  return $Self->{outpos};
}

sub ReadBlock {
  my $Self = shift;
  unless ($Self->{infh}) {
    die "Attempt to read without input filehandle";
  }
  my $bytes = BLOCKSIZE;
  my $buf;
  my @result;
  while ($bytes > 0) {
    my $n = sysread($Self->{infh}, $buf, $bytes);
    unless ($n) {
      delete $Self->{infh};
      return if ($bytes == BLOCKSIZE); # end of file on block boundary
      die "Failed to read full block at $Self->{inpos}";
    }
    $bytes -= $n;
    $Self->{inpos} += $n;
    push @result, $buf;
  }
  return join('', @result);
}

sub WriteBlock {
  my $Self = shift;
  my $string = shift;

  unless ($Self->{outfh}) {
    die "Attempt to write without output filehandle";
  }
  my $pos = $Self->{outpos};

  # make sure we've got BLOCKSIZE bytes to write
  if (length($string) > BLOCKSIZE) {
    $string = substr($string, 0, BLOCKSIZE);
  }
  elsif (length($string) < BLOCKSIZE) {
    $string .= "\0" x (BLOCKSIZE - length($string));
  }

  my $bytes = BLOCKSIZE;

  while ($bytes > 0) {
    my $n = syswrite($Self->{outfh}, $string, $bytes, BLOCKSIZE - $bytes);
    unless ($n) {
      delete $Self->{outfh};
      die "Failed to write full block at $Self->{outpos}";
    }
    $bytes -= $n;
    $Self->{outpos} += $n;
  }

  return $pos;
}

sub ReadHeader {
  my $Self = shift;
  my %Opts = @_;

  my ($pos, $header, $skipped) = (0, undef, 0);

  my $initialpos = $Self->{inpos};
  while (not $header) {
    $pos = $Self->{inpos};
    my $block = $Self->ReadBlock();
    last unless $block;
    $header = $Self->ParseHeader($block);
    last if $header;
    last unless $Opts{SkipInvalid};
    $skipped++;
  } 

  return unless $header;

  if ($skipped) {
    warn "Skipped $skipped blocks - invalid headers at $initialpos\n";
  }

  $header->{_pos} = $pos;
  $Self->{last_header} = $header;

  return $header;
}

sub WriteHeader {
  my $Self = shift;
  my $header = shift;

  my $block = $Self->CreateHeader($header);
  my $pos = $Self->WriteBlock($block);
  return( {%$header, _pos => $pos} );
}

sub ParseHeader {
  my $Self = shift;
  my $block = shift;

  # enforce length
  return unless(512 == length($block));

  # skip empty blocks
  return if substr($block, 0, 1) eq "\0";

  # unpack exactly 15 items from the block
  my @items = unpack("a100a8a8a8a12a12a8a1a100a8a32a32a8a8a155", $block);
  return unless (15 == @items);

  for (@items) {
    s/\0.*//; # strip from first null
  }

  my $chksum = oct($items[6]);
  # do checksum
  substr($block, 148, 8) = "        ";
  unless (unpack("%16C*", $block) == $chksum) {
    return;
  }

  my %header = (
    name => $items[0],
    mode => oct($items[1]),
    uid => oct($items[2]),
    gid => oct($items[3]),
    size => oct($items[4]),
    mtime => oct($items[5]),
    # checksum
    typeflag => $items[7],
    linkname => $items[8],
    # magic
    uname => $items[10],
    gname => $items[11],
    devmajor => oct($items[12]),
    devminor => oct($items[13]),
    prefix => $items[14],
  );

  return \%header;
}

sub BlankHeader {
  my $Self = shift;
  my %hash = (
    name => '',
    mode => 0777,
    uid => 0,
    gid => 0,
    size => 0,
    mtime => time(),
    typeflag => '0', # this is actually the STANDARD plain file format, phooey.  Not 'f' like Tar writes
    linkname => '',
    uname => '',
    gname => '',
    devmajor => 0,
    devminor => 0,
    prefix => '',
  );
  my %overrides = @_;
  foreach my $key (keys %overrides) {
    if (exists $hash{$key}) {
      $hash{$key} = $overrides{$key};
    }
    else {
      warn "invalid key $key for tar header\n";
    }
  }
  return \%hash;
}

sub CreateHeader {
  my $Self = shift;
  my $header = shift;

  my $block = pack("a100a8a8a8a12a12a8a1a100a8a32a32a8a8a155",
    $header->{name},
    sprintf("%07o", $header->{mode}),
    sprintf("%07o", $header->{uid}),
    sprintf("%07o", $header->{gid}),
    sprintf("%011o", $header->{size}),
    sprintf("%011o", $header->{mtime}),
    "        ",                  # chksum
    $header->{typeflag},
    $header->{linkname},
    "ustar  \0",                   # magic
    $header->{uname},
    $header->{gname},
    sprintf("%07o", $header->{devmajor}),
    sprintf("%07o", $header->{devminor}),
    $header->{prefix},
  );

  # calculate checksum
  my $checksum = sprintf("%06o", unpack("%16C*", $block));
  substr($block, 148, 8) = $checksum . "\0 ";

  # pad out to BLOCKSIZE characters
  if (length($block) < BLOCKSIZE) {
    $block .= "\0" x (BLOCKSIZE - length($block));
  }
  elsif (length($block) > BLOCKSIZE) {
    $block = substr($block, 0, BLOCKSIZE);
  }

  return $block;
}

# Chooser interface:
#
# $Chooser->($header, $outpos, undef);
# returns two values:
# ($status, $newheader)
#
# if $newheader is specified, it will be used instead.
#
# $status can be:
#   EXIT - quit streaming at this point, ignore this item
#   KEEP - copy this file as is (possibly changed header) to output tar
#   EDIT - re-call $Chooser with filehandle
#   SKIP - skip over the file and call $Chooser on the next one
#
# EDIT mode:
# the file will be copied to a temporary file and the filehandle passed to
# $Chooser.  It can truncate, rewrite, edit - whatever.  So long as it updates
# $header->{size} and returns it as $newheader it's all good.
#
# you don't have to change the file of course, it's also good just as a way to
# view the contents of some files as you stream them
sub StreamCopy {
  my $Self = shift;
  my $Chooser = shift;

  while (my $header = $Self->ReadHeader()) {
    my $pos = $header->{_pos};
    if ($Chooser) {
      my ($rc, $newheader) = $Chooser->($header, $Self->{outpos}, undef);
      
      # short circuit exit code
      return if $rc eq 'EXIT';

      # positive code means read the file
      if ($rc eq 'KEEP' or $rc eq 'EDIT') {
        my $TempFile = $Self->CopyToTempFile($header->{size});

        if ($rc eq 'EDIT') { 
          # call chooser again with the contents
          ($rc, $newheader) = $Chooser->($header, $Self->{outpos}, $TempFile);
          seek($TempFile, 0, 0);
        }

        # NOTE: even the size could have been changed if it's an edit!
        $header = $newheader if $newheader;

        if ($rc eq 'KEEP') {
          print "KEEP $header->{name} $pos/$Self->{outpos}\n" if $VERBOSE;
          $Self->WriteHeader($header);
          $Self->CopyFromFh($TempFile, $header->{size});
        }
        else {
          print "LATE REJECT $header->{name} $pos/$Self->{outpos}\n" if $VERBOSE;
        }
      }

      # anything else means discard it
      elsif ($rc eq 'SKIP') {
        print "DISCARD $header->{name} $pos/$Self->{outpos}\n" if $VERBOSE;
        $Self->DumpBytes($header->{size});
      }
      else {
        die "Bogus response $rc from callback\n";
      }
    }
    else {
      print "PASSTHROUGH $header->{name} $Self->{outpos}\n" if $VERBOSE;
      # XXX - faster but less safe
      #$Self->WriteHeader($header);
      #$Self->CopyBytes($header->{size});

      # slow safe option :)
      my $TempFile = $Self->CopyToTempFile($header->{size});
      $Self->WriteHeader($header);
      $Self->CopyFromFh($TempFile, $header->{size});
    }
  }
}

sub AddFile {
  my $Self = shift;
  my $name = shift;
  my $size = shift;
  my $fh = shift;

  my $header = $Self->BlankHeader(@_);
  $header->{name} = $name;
  $header->{size} = $size;

  my $fullheader = $Self->WriteHeader($header);
  $Self->CopyFromFh($fh, $size);

  return $fullheader;
}

sub AddLink {
  my $Self = shift;
  my $name = shift;
  my $linkname = shift;

  my $header = $Self->BlankHeader(@_);
  $header->{name} = $name;
  $header->{linkname} = $linkname;

  return $Self->WriteHeader($header);
}

sub CopyBytes {
  my $Self = shift;
  my $bytes = shift;
  while ($bytes > 0) {
    my $dump = $Self->ReadBlock();
    $Self->WriteBlock($dump);
    $bytes -= length($dump);
  }
}

sub DumpBytes {
  my $Self = shift;
  my $bytes = shift;
  while ($bytes > 0) {
    my $dump = $Self->ReadBlock();
    $bytes -= length($dump);
  }
}

sub FinishTar {
  my $Self = shift;
  $Self->WriteBlock("\0" x 512);
  $Self->WriteBlock("\0" x 512);
  $Self->WriteBlock("\0" x 512);
  $Self->WriteBlock("\0" x 512);
  $Self->WriteBlock("\0" x 512);
}

sub CopyToTempFile {
  my $Self = shift;
  my $bytes = shift;

  my $TempFile = File::Temp->new();
  while ($bytes > 0) {
    my $dump = $Self->ReadBlock();
    $dump = substr($dump, 0, $bytes) if length($dump) > $bytes;
    $TempFile->print($dump);
    $bytes -= length($dump);
  }
  seek($TempFile, 0, 0);

  return $TempFile;
}

sub CopyFromFh {
  my $Self = shift;
  my $Fh = shift;
  my $bytes = shift;

  my $buf;
  while ($bytes > 0) {
    my $thistime = $bytes > BLOCKSIZE ? BLOCKSIZE : $bytes;
    my $block = '';
    while ($thistime) {
      my $n = sysread($Fh, $buf, $thistime);
      unless ($n) {
        die "Failed to read entire file, doh ($bytes remaining)!\n";
      }
      $thistime -= $n;
      $block .= $buf;
    }
    if (length($block) < BLOCKSIZE) {
      $block .= "\0" x (BLOCKSIZE - length($block));
    }
    $Self->WriteBlock($block);
    $bytes -= BLOCKSIZE;
  }
}


# Finally,
1;
