#!/usr/bin/perl -w

##############################################
#
# aims2restoreimages.pl - Tool for dumping all pxe boot images to disk
#
# Author: Daniel Juarez Gonzalez <djuarezg@cern.ch>
#
##############################################

use strict;
use Getopt::Long;
use POSIX qw(setsid strftime);
use Sys::Syslog qw(:standard :macros);
use File::Path;
use IO::File;
use Digest::MD5;
use Date::Manip;
use Date::Parse;
use DBI;
use DBD::Oracle;
use File::Find::Rule;
use aims2server::conf;
use aims2server::db;
my ( $image_name, $destination_path, %code );

my $LOG = "/var/log/aims/aims2restoreimages.log";

##############################################
# Catch icky-things.
##############################################
$SIG{'INT'}   = \&interrupt;
$SIG{'HUP'}   = \&interrupt;
$SIG{'ABRT'}  = \&interrupt;
$SIG{'QUIT'}  = \&interrupt;
$SIG{'TRAP'}  = \&interrupt;
$SIG{'STOP'}  = \&interrupt;
$SIG{'TERM'}  = \&interrupt;
$SIG{__DIE__} = \&interrupt;

# Dump images from DB to disk; allow changing destination and selecting specific images

##############################################
my %opts = (
##############################################
    'image_name=s'       => \$image_name,
    'destination_path=s' => \$destination_path,
);
my $options = GetOptions(%opts);

##############################################
$code{'getimage'} = sub()
##############################################
{
    unless ($destination_path) {
        die "Error: Destination path required. Use aims2restoreimages getimage /destination image_name\n";
    }

    unless ($image_name) {
        die "Error: Image name required.\n";
    }

    my $dbh = db_connect();

    build_pxeboot( $dbh, $image_name );

    print "Image $image_name stored in $destination_path.\n";

};

##############################################
$code{'getallimages'} = sub()
##############################################
{
    unless ($destination_path) {
        die "Error: Destination path required. Use aims2restoreimages getallimages /destination\n";
    }

    my $dbh = db_connect();

    ##############################################
    # Get a list of pxeboot.
    ##############################################
    my $pxeboot = get_pxeboot_list($dbh);
    my $nowtime = strftime "%s", localtime();

    ##############################################
    # Remove junk from /boot/
    ##############################################
    remove_dead_pxeboot($pxeboot);    # remove unwanted crap.

    ##############################################
    # Build new pxeboot media.
    ##############################################
    add_new_pxeboot( $dbh, $pxeboot );

    foreach my $name ( keys %$pxeboot ) {
        build_pxeboot( $dbh, $name );
    }

    print "All images stored in $destination_path.\n";

};

##############################################
sub add_new_pxeboot
##############################################
{
    my ( $dbh, $pxeboot ) = @_;
    opendir( *pxedir, "$destination_path" )
      or die("cannot read contents of destination directory: $!");
    my @contents = grep m/^\w+?/ig, readdir *pxedir; # FIXME: warning, no regex.
    closedir(*pxedir);
    my %content = map  { $_, 1 } (@contents);
    my @addons  = map  { $_ } ( keys %$pxeboot );
    my @addon   = grep { !$content{$_} } @addons;
    foreach my $name (@addon) {
        build_pxeboot( $dbh, $name );
    }

}

##############################################
sub remove_dead_pxeboot
##############################################
{
    my ($pxeboot) = @_;
    opendir( *pxedir, "$destination_path" )
      or die("cannot read contents of TFTP_BOOT directory: $!");
    my @contents = grep m/^\w+?/ig, readdir *pxedir; # FIXME: warning, no regex.
    closedir(*pxedir);
    my %removals = map  { $_, 1 } ( keys %$pxeboot );
    my @removal  = grep { !$removals{$_} } @contents;
    foreach (@removal) {
        alog( $LOG, "removing pxeboot $_" );
        my $file = "$destination_path/$_";
        rmtree $file or groan("failed to remove $file: $!");
    }
}

##############################################
sub download_pxeboot
##############################################
{
    my ( $dbh, $name ) = @_;

    # Get pxeboot data.
    my @fields = qw(name vmlinuz_sum initrd_sum vmlinuz_file initrd_file);
    my $fields = join( ",", @fields );
    my $sql    = "SELECT $fields FROM pxeboot WHERE name = ?";
    my %files  = ();
    eval {
        my $prepare = $dbh->prepare($sql)
          or die("failed to prepare download of $name");
        $prepare->execute($name);
        $prepare->bind_columns( undef, map { \$_ } @fields );
        while ( my ( $NAME, @files ) = $prepare->fetchrow_array ) {
            (
                $files{$NAME}{VMLINUZ_SUM},  $files{$NAME}{INITRD_SUM},
                $files{$NAME}{VMLINUZ_BLOB}, $files{$NAME}{INITRD_BLOB},
            ) = @files;
        }
    };
    if ($@) {
        if ( $@ =~ $dbh->err ) {
            die("Failed download of pxeboot $name. see logs");
            alog( $LOG, "BLOB download failed for $name: $@" );
        }
    }
    \%files;
}

##############################################
sub get_file_sum
##############################################
{
    # Get an md5sum from the file.
    my $file   = shift;
    my $handle = new IO::File( $file, "r" );
    my $ident  = Digest::MD5->new;
    while (<$handle>) {
        $ident->add($_);
    }
    my $sum = $ident->hexdigest;
    return $sum;
}

##############################################
sub write_file
##############################################
{
    my ( $file, $data, $sum ) = @_;

    my $built = 0;

    # We may have already written the file ok the first time. Check.
    if ( -e $file ) {
        my $existing = get_file_sum($file);
        if ( $existing eq $sum ) {
            alog( $LOG, "skipping $file, already built" );
            $built = 1;
        }
    }

    # If we haven't.
    unless ( $built > 0 ) {
        my $build = new IO::File ">$file";
        if ( defined $build ) {
            print $build $data;
            $build->close;
        }
        else {
            die("cannot write file $file : $!");
        }

        # Check the file built correctly.
        my $SUM = get_file_sum($file);

        unless ( $SUM eq $sum ) {
            die("file mismatch with $file ($SUM <> $sum)");
        }
    }
}

##############################################
sub build_pxeboot
##############################################
{
    my ( $dbh, $NAME ) = @_;

    # Build the files.

    my $files = download_pxeboot( $dbh, $NAME );
    alog( $LOG, "downloaded $NAME" );
    my $dir = "$destination_path/$NAME";
    unless ( -e $dir ) {
        mkdir $dir or groan("Cannot create directory for $NAME");
    }

    # Create the kernel.
    write_file(
        "$dir/vmlinuz",
        $files->{$NAME}{VMLINUZ_BLOB},
        $files->{$NAME}{VMLINUZ_SUM}
    );
    alog( $LOG, "built vmlinuz for $NAME" );

    # Create initrd.img (if present)
    if ( $files->{$NAME}{INITRD_SUM} ) {
        write_file(
            "$dir/initrd",
            $files->{$NAME}{INITRD_BLOB},
            $files->{$NAME}{INITRD_SUM}
        );
        alog( $LOG, "built initrd.img for $NAME" );
    }

}

##############################################
sub groan
##############################################
{
    my $groan = shift;
    alog( $LOG, "error: $groan" );
}

##############################################
sub interrupt {
##############################################
    # Handle things nicely when we croak.
    alog( $LOG, "caught @_, exiting." );
    printf STDOUT "caught @_ exiting\n";

    # Self killing this script: https://perldoc.perl.org/perlvar#%24%24
    # Cannot exit or die, it requires ctrl+c per image, just kill itself
    kill 9, -$$;
}

##############################################
sub db_connect
##############################################
{
    my $dbconn =
        $aims2conf::db_conn_string
      ? $aims2conf::db_conn_string
      : die "Error: Database server undefined.\n";
    my $dbuser =
        $aims2conf::db_user
      ? $aims2conf::db_user
      : die "Error: Database user undefined.\n";
    my $dbpass =
        $aims2conf::db_pass
      ? $aims2conf::db_pass
      : die "Error: Database password undefined.\n";
    my %dbattrs =
        %aims2conf::db_conn_attrib
      ? %aims2conf::db_conn_attrib
      : die "Error: Database attributes undefined.\n";

    my $db  = aims2db->new();
    my $dbh = $db->Connect( $dbconn, $dbuser, $dbpass, %dbattrs );

    unless ($dbh) {
        die "Error: Cannot establish database connection. Exiting...\n";
    }
    $dbh;
}

##############################################
sub get_pxeboot_list
##############################################
{
    my $dbh     = shift;
    my %pxeboot = ();
    my $sql =
      "SELECT name,to_char(uploaded, 'yyyy/mm/dd hh24:mi:ss') FROM pxeboot";
    eval {
        my $prepare = $dbh->prepare($sql)
          or die( "failed to prepare pxeboot query: " . $dbh->errstr );
        $prepare->execute();
        while ( my ( $NAME, @meta ) = $prepare->fetchrow_array ) {
            (
                $pxeboot{$NAME}{UPLOADED},

            ) = @meta;
        }
    };
    if ($@) {
        if ( $@ =~ $dbh->err ) {
            die("failed execution of pxeboot list query. see logs.");
            alog( $LOG, "query failed: $@" );
        }
    }
    \%pxeboot;
}

##############################################
sub insert0 {
##############################################
    # 7th July =~ 07th July :)
    my ($date) = shift;
    if ( $date < 10 ) {
        return "0$date";
    }
    return $date;
}

##############################################
sub getdate {
##############################################
    my ( $sec, $min, $hour, $mday, $mon, $year, $wday, $yday, $iddst ) =
      localtime(time);
    my $datestring;
    $year += 1900;
    $mon++;
    $mon        = insert0($mon);
    $mday       = insert0($mday);
    $min        = insert0($min);
    $datestring = "$year-$mon-$mday $hour:$min";
    return ($datestring);
}

##############################################
sub alog {
##############################################
    my ( $logfile, $status ) = @_;

    # FIXME: This is the date it was logged, not executed.. naughty.
    my $date = getdate();
    if ( open( FILE, ">>$logfile" ) ) {
        print FILE ("$date $status\n");
        close FILE;
    }

    syslog( "info", $status . "\n" );
}

##############################################
sub init
##############################################
{

    if ( !@ARGV ) {
        print
"Usage: $0 getimage|getallimages| [--destination_path|--image_name]\n";
        exit 1;
    }

    my $cmd = $ARGV[0];
    unless ($cmd) {
        exit 0;
    }

    $destination_path = $ARGV[1] ? $ARGV[1] : undef;

    $image_name = $ARGV[2];
    $image_name =~ tr/[a-z]/[A-Z]/ if $image_name;

    if ( $code{$cmd} ) {
        print
          "Running, check the logs on /var/log/aims/aims2restoreimages.log\n";
        $code{$cmd}->();    # Go!
    }
    else {
        exit 1;
    }
}

init();
