#!/usr/bin/perl -w
#
# Copyright (c) 2018 SUSE Inc.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 as
# published by the Free Software Foundation.
#
# 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 (see the file COPYING); if not, write to the
# Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA
#
################################################################
#
# Registry interfacing
#

BEGIN {
  my ($wd) = $0 =~ m-(.*)/- ;
  $wd ||= '.';
  unshift @INC,  "$wd/build";
  unshift @INC,  "$wd";
}

use JSON::XS ();
use Digest::SHA ();
use Data::Dumper;
use Compress::Zlib;

use BSRPC ':https';
use BSHTTP;
use BSTar;
use BSContar;
use BSUtil;
use BSBearer;

use strict;

my $registry_timeout = 300;

my $dest_creds;
my $use_image_tags;
my $multiarch;
my $digestfile;
my $delete_mode;
my $delete_except_mode;
my $list_mode;
my @tags;
my $blobdir;

my $registryserver;
my $repository;
my @tarfiles;

my $registry_authenticator;

sub send_layer {
  my ($param) = @_;
  my $sl = $param->{'send_layer_data'}; # [ $layer_ent, $offset ],
  my $chunk = BSTar::extract($sl->[0]->{'file'}, $sl->[0], $sl->[1], 65536);
  $sl->[1] += length($chunk);
  return $chunk;
}

sub blob_exists {
  my ($blobid, $size) = @_;
  my $replyheaders;
  my $param = {
    'uri' => "$registryserver/v2/$repository/blobs/$blobid",
    'request' => 'HEAD',
    'authenticator' => $registry_authenticator,
    'timeout' => $registry_timeout,
    'replyheaders' => \$replyheaders,
    'maxredirects' => 1,
  };
  eval { BSRPC::rpc($param) };
  if ($replyheaders) {
    die("size mismatch?\n") if $replyheaders->{'content-length'} && $size && $size != $replyheaders->{'content-length'};
    return 0 if $replyheaders->{'docker-content-digest'} && $blobid ne $replyheaders->{'docker-content-digest'};
    return 1;
  }
  return 0;
}

sub blob_upload {
  my ($blobid, $upload_ent) = @_;

  return $blobid if blob_exists($blobid, $upload_ent->{'size'});
  print "uploading layer $blobid... ";
  my $replyheaders;
  my $param = {
    'uri' => "$registryserver/v2/$repository/blobs/uploads/",
    'request' => 'POST',
    'authenticator' => $registry_authenticator,
    'timeout' => $registry_timeout,
    'replyheaders' => \$replyheaders,
  };
  BSRPC::rpc($param);
  my $loc = $replyheaders->{'location'};
  my @locextra;
  if ($loc =~ s/\?(.*)$//) {
    push @locextra, split('&', $1);
    s/%([a-fA-F0-9]{2})/chr(hex($1))/ge for @locextra;
  }
  die("no location in upload reply\n") unless $loc;
  $loc = "$registryserver$loc" if $loc =~ /^\//;
  $param = {
    'headers' => [ "Content-Length: $upload_ent->{'size'}", "Content-Type: application/octet-stream" ],
    'uri' => $loc,
    'request' => 'PUT',
    'authenticator' => $registry_authenticator,
    'replyheaders' => \$replyheaders,
    'data' => \&send_layer,
    'send_layer_data' => [ $upload_ent, 0 ],
  };
  $replyheaders = undef;
  BSRPC::rpc($param, undef, @locextra, "digest=$blobid");
  my $id = $replyheaders->{'docker-content-digest'};
  die("server did not return a content id\n") unless $id;
  die("server created a different blobid: $blobid $id\n") if $blobid ne $id;
  print "done.\n";
  return $blobid;
}

sub blob_download {
  my ($blobid, $filename) = @_;
  my $stdout_receiver = sub {
    while(1) {
      my $s = BSHTTP::read_data($_[0], 8192);
      return {} if $s eq '';
      print($s) || die("write: $!\n");
    }
  };
  my $param = {
    'uri' => "$registryserver/v2/$repository/blobs/$blobid",
    'authenticator' => $registry_authenticator,
    'timeout' => $registry_timeout,
    'receiver' => $filename eq '-' ? $stdout_receiver : \&BSHTTP::file_receiver,
    'filename' => $filename,
    'maxredirects' => 1,
  };
  BSRPC::rpc($param);
}

sub blob_fetch {
  my ($blobid) = @_;
  my $param = {
    'uri' => "$registryserver/v2/$repository/blobs/$blobid",
    'authenticator' => $registry_authenticator,
    'timeout' => $registry_timeout,
    'maxredirects' => 1,
  };
  return BSRPC::rpc($param);
}

sub manifest_exists {
  my ($manifest, $tag, $content_type) = @_;

  $content_type ||= 'application/vnd.docker.distribution.manifest.v2+json';
  my $maniid = 'sha256:'.Digest::SHA::sha256_hex($manifest);
  $tag = $maniid unless defined $tag;
  my $replyheaders;
  my $param = {
    'headers' => [ "Accept: $content_type" ],
    'uri' => "$registryserver/v2/$repository/manifests/$tag",
    'authenticator' => $registry_authenticator,
    'replyheaders' => \$replyheaders,
    'timeout' => $registry_timeout,
  };
  my $maniret;
  eval { $maniret = BSRPC::rpc($param) };
  if ($maniret) {
    my $maniretid = $replyheaders->{'docker-content-digest'};
    $maniretid ||= 'sha256:'.Digest::SHA::sha256_hex($maniret);
    return 1 if $maniretid eq $maniid;
  }
  return 0;
}

sub manifest_upload {
  my ($manifest, $tag, $content_type, $quiet) = @_;

  my $maniid = 'sha256:'.Digest::SHA::sha256_hex($manifest);
  return $maniid if manifest_exists($manifest, $tag, $content_type);
  $content_type ||= 'application/vnd.docker.distribution.manifest.v2+json';
  my $fat = '';
  $fat = 'fat ' if $content_type eq 'application/vnd.docker.distribution.manifest.list.v2+json';
  if (!$quiet) {
    if (defined($tag)) {
      print "uploading ${fat}manifest $maniid for tag '$tag'... ";
    } else {
      print "uploading ${fat}manifest $maniid... ";
    }
  }
  $tag = $maniid unless defined $tag;
  my $replyheaders;
  my $param = {
    'headers' => [ "Content-Type: $content_type" ],
    'uri' => "$registryserver/v2/$repository/manifests/$tag",
    'request' => 'PUT',
    'authenticator' => $registry_authenticator,
    'replyheaders' => \$replyheaders,
    'data' => $manifest,
  };
  BSRPC::rpc($param, undef);
  my $id = $replyheaders->{'docker-content-digest'};
  die("server did not return a content id\n") unless $id;
  die("server created a different manifest id: $maniid $id\n") if $maniid ne $id;
  print "done.\n" unless $quiet;
  return $maniid;
}

sub manifest_append {
  my ($manifest, $tag) = @_;
  my $maniid = 'sha256:'.Digest::SHA::sha256_hex($manifest);
  my $str = "$maniid ".length($manifest)." $tag\n";
  if ($digestfile ne '-') {
    BSUtil::appendstr($digestfile, $str);
  } else {
    print $str;
  }
}

sub manifest_upload_tags {
  my ($manifest, $tags, $content_type) = @_;
  return manifest_upload($manifest, undef, $content_type) unless @{$tags || []};
  for my $tag (BSUtil::unify(@$tags)) {
    manifest_upload($manifest, $tag, $content_type);
    manifest_append($manifest, $tag) if defined $digestfile;
  }
}

sub get_all_tags {
  my @regtags;
  my $replyheaders;
  my $param = {
    'uri' => "$registryserver/v2/$repository/tags/list",
    'authenticator' => $registry_authenticator,
    'timeout' => $registry_timeout,
    'replyheaders' => \$replyheaders,
  };
  while (1) {
    undef $replyheaders;
    my $r = BSRPC::rpc($param, \&JSON::XS::decode_json);
    push @regtags, @{$r->{'tags'} || []};
    last unless $replyheaders->{'link'};
    die unless $replyheaders->{'link'} =~ /^<(\/v2\/.*)>/;
    $param->{'uri'} = "$registryserver$1";
    $param->{'verbatim_uri'} = 1;
  }
  return BSUtil::unify(@regtags);
}

sub get_all_repositories {
  my @repos;
  my $replyheaders;
  my $param = {
    'uri' => "$registryserver/v2/_catalog",
    'authenticator' => $registry_authenticator,
    'timeout' => $registry_timeout,
    'replyheaders' => \$replyheaders,
  };
  while (1) {
    undef $replyheaders;
    my $r = BSRPC::rpc($param, \&JSON::XS::decode_json);
    push @repos, @{$r->{'repositories'}};
    last unless $replyheaders->{'link'};
    die unless $replyheaders->{'link'} =~ /^<(\/v2\/.*)>/;
    $param->{'uri'} = "$registryserver$1";
    $param->{'verbatim_uri'} = 1;
  }
  return BSUtil::unify(@repos);
}

sub get_manifest_for_tag {
  my ($tag) = @_;
  my $replyheaders;
  my $param = {
    'uri' => "$registryserver/v2/$repository/manifests/$tag",
    'headers' => [ "Accept: application/vnd.docker.distribution.manifest.list.v2+json, application/vnd.docker.distribution.manifest.v2+json" ],
    'authenticator' => $registry_authenticator,
    'replyheaders' => \$replyheaders,
    'timeout' => $registry_timeout,
  };
  my $mani_json;
  eval { $mani_json = BSRPC::rpc($param); };
  if ($@) {
    return () if $@ =~ /^404/;	# tag does not exist
    die($@);
  }
  my $mani = JSON::XS::decode_json($mani_json);
  my $maniid = $replyheaders->{'docker-content-digest'};
  die("$tag: no docker-content-digest\n") unless $maniid;
  return ($mani, $maniid, $mani_json);
}

sub delete_tag {
  my ($tag) = @_;

  my ($mani, $maniid, $mani_json) = get_manifest_for_tag($tag);
  return unless $maniid;
  print "deleting tag $tag [$maniid]\n";

  # now mangle and upload so that we get a new unique image id for that tag
  my $mani_json_mangled = "\n\n".JSON::XS->new->utf8->canonical->encode($mani)."\n\n";
  my $newmaniid = manifest_upload($mani_json_mangled, $tag, $mani->{'mediaType'}, 1);

  # then delete the new unique image id, thus deleting the tag as well
  my $param = {
    'uri' => "$registryserver/v2/$repository/manifests/$newmaniid",
    'request' => 'DELETE',
    'authenticator' => $registry_authenticator,
    'timeout' => $registry_timeout,
  };
  eval { BSRPC::rpc($param, undef); };
  my $err = $@;

  if ($err) {
    # delete failed, switch tag back to the old image
    eval { manifest_upload($mani_json, $tag, $mani->{'mediaType'}, 1) if $newmaniid ne $maniid; };
    die($err);
  }
}

sub list_tag {
  my ($tag) = @_;
  
  my ($mani, $maniid) = get_manifest_for_tag($tag);
  if (!$mani) {
    printf "%-20s -\n", $tag;
  } elsif (!$mani->{'mediaType'}) {
    printf "%-20s %s %s\n", $tag, $maniid, 'v1image';
  } elsif ($mani->{'mediaType'} eq 'application/vnd.docker.distribution.manifest.list.v2+json') {
    printf "%-20s %s %s\n", $tag, $maniid, 'list';
  } elsif ($mani->{'mediaType'} eq 'application/vnd.docker.distribution.manifest.v2+json') {
    printf "%-20s %s %s\n", $tag, $maniid, 'image';
  } else {
    printf "%-20s %s %s\n", $tag, $maniid, 'unknown';
  }
}

sub tags_from_digestfile {
  return () unless $digestfile;
  my @ret;
  local *DIG;
  open(DIG, '<', $digestfile) || die("$digestfile: $!\n");
  while (<DIG>) {
    chomp;
    next if /^#/ || /^\s*$/;
    die("bad line in digest file\n") unless /^([a-z0-9]+):([a-f0-9]+) (\d+) (.+?)\s*$/;
    push @ret, $4;
  }
  close(DIG);
  return @ret;
}

sub construct_container_tar {
  my ($containerinfo) = @_; 

  die("Must specify a blobdir for containerinfos\n") unless $blobdir;
  my $manifest = $containerinfo->{'tar_manifest'};
  my $mtime = $containerinfo->{'tar_mtime'};
  my $blobids = $containerinfo->{'tar_blobids'};
  die("containerinfo is incomplete\n") unless $mtime && $manifest && $blobids;
  my @tar;
  for my $blobid (@$blobids) {
    my $fd;
    open($fd, '<', "$blobdir/_blob.$blobid") || die("$blobdir/_blob.$blobid: $!\n");
    push @tar, {'name' => $blobid, 'file' => $fd, 'mtime' => $mtime, 'offset' => 0, 'size' => (-s $fd)};
  }
  push @tar, {'name' => 'manifest.json', 'data' => $manifest, 'mtime' => $mtime, 'size' => length($manifest)};
  return \@tar;
}

sub open_tarfile {
  my ($tarfile) = @_;

  my ($tar, $tarfd, $variant);
  if ($tarfile =~ /\.containerinfo$/) {
    my $containerinfo_json = readstr($tarfile);
    my $containerinfo = JSON::XS::decode_json($containerinfo_json);
    $variant = $containerinfo->{'govariant'};
    $tar = construct_container_tar($containerinfo);
  } else {
    open($tarfd, '<', $tarfile) || die("$tarfile: $!\n");
    $tar = BSTar::list($tarfd);
    $_->{'file'} = $tarfd for @$tar;
  }
  return ($tar, $tarfd, $variant);
}

sub die_with_usage {
  die <<'END';
usage: bs_regpush [options] <registryserver> repository tar [tar...]
       bs_regpush [options] -l <registryserver> [repository [tag]]
       bs_regpush [options] -D <registryserver> repository
       bs_regpush [options] -X <registryserver> repository

MODES:
                - upload mode
  -l            - list mode
  -D            - delete mode
  -X            - delete except mode

OPTIONS:
  --dest-creds  - credentials in form <user>:<password> or "-" to read from STDIN
  -T            - use image tags
  -m            - push multiarch image
  -t            - tag (can be given multiple times)
  -F            - digestfile, output in upload mode, otherwise input

END
}

$| = 1;

while (@ARGV) {
  if ($ARGV[0] eq '--dest-creds') {
    $dest_creds = BSBearer::get_credentials($ARGV[1]);
    splice(@ARGV, 0, 2);
  } elsif ($ARGV[0] eq '-T') {
    $use_image_tags = 1;
    shift @ARGV;
  } elsif ($ARGV[0] eq '-m') {
    $multiarch = 1;
    shift @ARGV;
  } elsif ($ARGV[0] eq '-t') {
    push @tags, $ARGV[1];
    splice(@ARGV, 0, 2);
  } elsif ($ARGV[0] eq '-F') {
    $digestfile = $ARGV[1];
    splice(@ARGV, 0, 2);
  } elsif ($ARGV[0] eq '-D') {
    $delete_mode = 1;
    shift @ARGV;
  } elsif ($ARGV[0] eq '-l') {
    $list_mode = 1;
    shift @ARGV;
  } elsif ($ARGV[0] eq '-X') {
    $delete_except_mode = 1;
    shift @ARGV;
  } elsif ($ARGV[0] eq '-B') {
    $blobdir = $ARGV[1];
    splice(@ARGV, 0, 2);
  } else {
    last;
  }
}

$registry_authenticator = BSBearer::generate_authenticator($dest_creds, 'verbose' => (-c STDOUT ? 1 : 0));

if ($list_mode) {
  ($registryserver, $repository) = @ARGV;
  if (@ARGV == 1) {
    for my $repo (sort(get_all_repositories())) {
      print "$repo\n";
    }
  } elsif (@ARGV == 2) {
    my %tags = map {$_ => 1} @tags;
    $tags{$_} = 1 for tags_from_digestfile();
    %tags = map {$_ => 1} get_all_tags() unless %tags;
    list_tag($_) for sort keys %tags;
  } elsif (@ARGV == 3) {
    my ($mani, $maniid, $mani_json) = get_manifest_for_tag($ARGV[2]);
    print "$mani_json\n" if $mani_json;
  } elsif (@ARGV == 4) {
    if ($ARGV[3] eq 'config' || $ARGV[3] eq 'config.json') {
      my ($mani) = get_manifest_for_tag($ARGV[2]);
      my $config = blob_fetch($mani->{'config'}->{'digest'});
      $config = JSON::XS->new->utf8->canonical->pretty->encode(JSON::XS::decode_json($config)) if $ARGV[3] eq 'config';
      print $config;
    } else {
      blob_download($ARGV[2], $ARGV[3]);
    }
  } else {
    die_with_usage();
  }
  exit(0);
}

die_with_usage() unless @ARGV >= 2;
($registryserver, $repository, @tarfiles) = @ARGV;

if ($delete_mode || $delete_except_mode) {
  die("cannot do both delete and delete-except\n") if $delete_mode && $delete_except_mode;
  my %tags;
  $tags{$_} = 1 for @tags;
  $tags{$_} = 1 for tags_from_digestfile();

  if ($delete_mode) {
    for my $tag (sort keys %tags) {
      delete_tag($tag);
    }
  } elsif ($delete_except_mode) {
    for my $tag (grep {!$tags{$_}} get_all_tags()) {
      delete_tag($tag);
    }
  }
  exit;
}

die("No tar file to upload?\n") if !@tarfiles;
die("more than one tar file specified\n") if @tarfiles > 1 && !$multiarch;

if ($use_image_tags && @tarfiles > 1) {
  # make sure all tar files contain the same tags
  my $imagetags;
  for my $tarfile (@tarfiles) {
    my ($tar, $tarfd) = open_tarfile($tarfile);
    my %tar = map {$_->{'name'} => $_} @$tar;
    my ($manifest_ent, $manifest) = BSContar::get_manifest(\%tar);
    my @imagetags = @{$manifest->{'RepoTags'} || []};
    s/.*:// for @imagetags;
    my $it = join(', ', sort(BSUtil::unify(@imagetags)));
    die("multiarch images contain different tags: $imagetags -- $it\n") if defined($imagetags) && $imagetags ne $it;
    if (!defined($imagetags)) {
      $imagetags = $it;
      push @tags, @imagetags;
    }
    close $tarfd if $tarfd;
  }
  $use_image_tags = 0;
}

my @multimanifests;
my %multiplatforms;
for my $tarfile (@tarfiles) {
  my ($tar, $tarfd, $variant) = open_tarfile($tarfile);
  my %tar = map {$_->{'name'} => $_} @$tar;

  my ($manifest_ent, $manifest) = BSContar::get_manifest(\%tar);
  #print Dumper($manifest);

  if ($use_image_tags) {
    my @imagetags = @{$manifest->{'RepoTags'} || []};
    s/.*:// for @imagetags;
    push @tags, @imagetags if $use_image_tags;
  }

  my ($config_ent, $config) = BSContar::get_config(\%tar, $manifest);
  #print Dumper($config);

  my @layers = @{$manifest->{'Layers'} || []};
  die("container has no layers\n") unless @layers;
  my $config_layers = $config->{'rootfs'}->{'diff_ids'};
  die("layer number mismatch\n") if @layers != @{$config_layers || []};

  $variant = $config->{'variant'} if $config->{'variant'};
  if ($multiarch) {
    # see if a already have this arch/os combination
    my $platformstr = "architecture:$config->{'architecture'} os:$config->{'os'}";
    $platformstr .= " variant:$variant" if $variant;
    if ($multiplatforms{$platformstr}) {
      print "ignoring $tarfile, already have $platformstr\n";
      close $tarfd if $tarfd;
      next;
    }
    $multiplatforms{$platformstr} = 1;
  }


  # process config
  my $config_blobid = BSContar::blobid_entry($config_ent);
  # create layer data
  my $config_data = { 
    'mediaType' => 'application/vnd.docker.container.image.v1+json',
    'size' => $config_ent->{'size'},
    'digest' => $config_blobid,
  }; 
  # upload to server
  blob_upload($config_blobid, $config_ent);
  
  # process layers (compress if necessary)
  my %layer_datas;
  my @layer_data;
  for my $layer_file (@layers) {
    if ($layer_datas{$layer_file}) {
      # already did that file, just reuse old layer data
      push @layer_data, $layer_datas{$layer_file};
      next;
    }
    my $layer_ent = $tar{$layer_file};
    die("File $layer_file not included in tar\n") unless $layer_ent;

    # detect layer compression
    my $comp = BSContar::detect_entry_compression($layer_ent);
    die("unsupported compression $comp\n") if $comp && $comp ne 'gzip';
    if (!$comp) {
      print "compressing $layer_ent->{'name'}... ";
      $layer_ent = BSContar::compress_entry($layer_ent);
      print "done.\n";
    }
    my $blobid = BSContar::blobid_entry($layer_ent);
    #print "$layer_file -> $blobid\n";

    # create layer data
    my $layer_data = {
      'mediaType' => 'application/vnd.docker.image.rootfs.diff.tar.gzip',
      'size' => $layer_ent->{'size'},
      'digest' => $blobid,
    };
    push @layer_data, $layer_data;
    $layer_datas{$layer_file} = $layer_data;

    # upload to server
    blob_upload($blobid, $layer_ent);
  }
  close $tarfd if $tarfd;

  my $mani = {
    'schemaVersion' => 2,
    'mediaType' => 'application/vnd.docker.distribution.manifest.v2+json',
    'config' => $config_data,
    'layers' => \@layer_data,
  };
  my $mani_json = BSContar::create_dist_manifest($mani);

  if ($multiarch) {
    manifest_upload($mani_json, undef);		# upload anonymous image
    my $multimani = {
      'mediaType' => 'application/vnd.docker.distribution.manifest.v2+json',
      'size' => length($mani_json),
      'digest' => 'sha256:'.Digest::SHA::sha256_hex($mani_json),
      'platform' => {'architecture' => $config->{'architecture'}, 'os' => $config->{'os'}},
    };
    $multimani->{'platform'}->{'variant'} = $variant if $variant;
    push @multimanifests, $multimani;
  } else {
    manifest_upload_tags($mani_json, \@tags);
  }
}

if ($multiarch) {
  my $ct = 'application/vnd.docker.distribution.manifest.list.v2+json';
  my $mani = {
    'schemaVersion' => 2,
    'mediaType' => $ct,
    'manifests' => \@multimanifests,
  };
  my $mani_json = BSContar::create_dist_manifest_list($mani);
  manifest_upload_tags($mani_json, \@tags, $ct);
}

