#!/usr/bin/env perl

# ###############################################
#
# packmule: make unattended loading environment
#
# Austin Shafer - 2018
#     ashaferian@gmail.com
#
# ###############################################

# ###############################################
# Copyright 2018-2019 Austin Shafer. All rights reserved
# 
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions
#  are met:
#  1. Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
#  2. Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
# 
#  THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
#  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
#  ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
#  FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
#  DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
#  OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
#  HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
#  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
#  OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
#  SUCH DAMAGE.
#
# ###############################################

# This program is designed to fill an freebsd installation image
# with any packages supplied by the user. This results in an image
# which holds extra utilities installed by the user.

# The script mounts the .iso image under /mnt/bsd_iso and copies
# the file system to /tmp/bsd-iso-work. Packages and software is
# added to under /tmp/bsd-iso-work and written to a new .iso image

use strict;

use YAML;
use Cwd;
use Digest::SHA;

# base prefix of installed location
my $PREFIX = "/usr/local/packmule/";

# global variables
my $mntdir = "/mnt/packmule-iso/";
my $rootprefix = "/tmp/packmule/";

# location of archdep scripts
my $imgprefix = "/mkiso/";
# location of image generation scripts
my $makeiso = "/mkisoimages.sh";
my $makeimg = "/make-memstick.sh";
# default to building iso
my $makescript = $makeiso;
# argument for mount
my $mount_arg = "-t cd9660";

# does /usr/src/release exist
my $release_exists = 1;

# directory holding copy of original iso
my $workroot = "$rootprefix-work.$$";
# directory where custom software is installed
my $pkgroot = "$rootprefix-pkgs.$$";
my $distdir = "$workroot/usr/freebsd-dist/";
my $md_dev = 4;
# Used for img (-U)
my $md_dev_slice = "";
# image extension
my $ext = "iso";

my $pkg = "export ASSUME_ALWAYS_YES=yes; pkg -r $pkgroot install";
my $distname = "packmule-pkgs.txz";

# the packing list config file
my $plist = "packmule.yml";
# the name of the image to generate
my $isoname;

# debug option
# will do everything except generate the iso
my $DEBUG;
my $REMOVE;

# ################################################
# Function declarations
# ################################################

# mounts isoname at mntdir
sub mount_img {
    print "Mounting $isoname at $mntdir\n";
    # NOTE:
    #       md node should be /dev/md${MD_DEV} for default
    # create vnode to mount under
    system("mdconfig -a -t vnode -f $isoname -u $md_dev");
    system("mount $mount_arg /dev/md$md_dev$md_dev_slice $mntdir");
}

# ################################################
# unmounts ISONAME at MNTDIR
sub unmount_img {
    print "Unmounting $isoname from $mntdir\n";

    system("umount $mntdir");
    # delete vnode (default is /dev/md${MD_DEV})
    system("mdconfig -d -u $md_dev");
}

# ################################################
# adds $distname to /usr/freebsd-dist/MANIFEST
sub manifest_add {
    my $distdir = "/usr/freebsd-dist/";

    open(my $dist, "<", "$workroot/$distdir/$distname");
    binmode($dist);

    # generate sha256 checksum
    my $hash = Digest::SHA->new(256)->addfile($dist)->hexdigest;
    close $dist;

    # calculate number of archived packages
    my $nfiles = `tar -t -f $workroot/$distdir/$distname | wc -l`;
    # trim whitespace
    $nfiles =~ s/[^0-9]//g;

    # append packmule-pkgs.txz to $distdir/MANIFEST
    open(my $manifest, ">>", "$workroot/$distdir/MANIFEST");

    # from /usr/src/release/scripts/make-manifest.sh
    print $manifest "$distname\t$hash\t$nfiles\tpackmule\t\"Packmule custom pkgs\"\ton\n";

    close $manifest;
}

# ################################################
# generates a new image (.iso or .img controlled by makescript)
sub make_img {
    my ($targ, $vol, $newiso) = @_;
    
    print "VOLUME_LABEL = $targ\n";

    # run machine dependent script to build install .iso
    if ($imgprefix eq "/mkiso") {
	system("sh $PREFIX/$imgprefix/$targ/$makescript -b $vol $newiso $workroot");
    } else {
	system("sh $PREFIX/$imgprefix/$targ/$makescript $workroot $newiso");
    }
}

# ################################################
# writes WORKROOT to a new .iso file
sub write_new_img {

    # change name from FreeBSD___.iso to FreeBSD___-packed.iso
    my $newisoname;
    if ($isoname =~ /.*\/([^\/]*).$ext|([^\/]*).$ext/) {
	$newisoname = "$1$2-packed.$ext";
    } else {
	$newisoname = "FreeBSD-packed.$ext";
    }

    print "Writing new install package ./$newisoname\n";

    my $target;
    # 4th group of something followed by '-'
    if ($newisoname =~ /(([^-]*)-){4}/) {
	$target = $2;
    }

    # construct volume label from isoname
    my $volume_label;
    if ($newisoname =~ /-(.*)-.*/) {
	$volume_label = $1;
	# change all dots and slashes to underscores
	$volume_label =~ s/(\.|-)/_/g;
    }

    make_img($target, $volume_label, $newisoname);
}

# ################################################
# A bit of a hack, but copying the password db is
# used because symlinks are not allowed. The passwd
# program will just create a new db overwriting the
# symlinks...

# ################################################
# copy password db so that
# installed packages can create groups
sub link_passwd {
    system("mkdir -p $pkgroot/etc/");

    system("cp $workroot/etc/passwd $pkgroot/etc/passwd");
    system("cp $workroot/etc/master.passwd $pkgroot/etc/master.passwd");
    system("cp $workroot/etc/pwd.db $pkgroot/etc/pwd.db");
    system("cp $workroot/etc/spwd.db $pkgroot/etc/spwd.db");
    system("cp $workroot/etc/group $pkgroot/etc/group");
    system("cp $workroot/etc/shells $pkgroot/etc/shells");
}

# ################################################
# removes directory listings
sub remove_dirs {
    system("rm -rf $rootprefix*");
}

# ################################################
# BEGIN MAIN
# ################################################

# check if /usr/src/release exits. If so, use the $makescript from
# the src tree instead of packmule's copies. This should catch
# new mkisoimages scripts if the users tree is populated with them
open "/usr/src/release/" or $release_exists = 0;

# release_exists is 1 if true
if ($release_exists) {
    $PREFIX = "/usr/src/release/";
    $imgprefix = "";
}

# #######################
# process command line arguments
while (@ARGV != 0) {
    my $arg = shift;

    if ($arg eq "-y") {
	$plist = shift;

    } elsif ($arg eq "-D") {
	$DEBUG = 1;

    } elsif ($arg eq "--no-pkg-scripts" or $arg eq "-I") {
	# do not run pre-install / post-install scripts for
	# compatibility sake. Some scripts fail due to the
	# '-r' option on pkg
	$pkg .= " -I ";

    } elsif ($arg eq "-R") {
	$REMOVE = 1;
	
    } elsif ($arg eq "-U") {
	# generate a USB img. This uses a different make script
	# for creating USB compatible partitions
	$makescript = $makeimg;
	$imgprefix = "/make-memstick";
	# The mount type needs to be chaned from cd9660, still ro
	$mount_arg = "-r";
	# the first slice is the efi section
	# we need the a partition of the second slice
	$md_dev_slice = "s2a";
	$ext = "img";

    } elsif (!defined($isoname)) {
	$isoname = $arg;

    } else {
	print "Skipping unrecognized argument $arg\n";
    }
}

# #######################
if (defined($REMOVE)) {
    print "packmule: Removing mounts.\n";
    remove_dirs();
    unmount_img();
    exit(0);
}
# isoname must be set
if (!defined($isoname)) {
    print "Usage: packmule [-y /path/to/config.yml] isoname\n";
    print "       (Default config is ./packmule.yaml)\n";
    exit(1);
}  

# parse pack list
my $config = YAML::LoadFile($plist);

# create directories if needed
remove_dirs;
system("mkdir -p $mntdir");
system("mkdir -p $workroot");
system("mkdir -p $pkgroot");

# mount the .iso
mount_img;

# copy image to r/w direcotry $workroot
print "--------------------------------------------------\n";
print "Copying filesystem from $mntdir to $workroot...\n";
print "--------------------------------------------------\n";
system("cp -R $mntdir/* $workroot/");

# use pkg to fill new image tree
print "--------------------------------------------------\n";
print "Installing packages in $pkgroot...\n";
print "--------------------------------------------------\n";

# setup

# create temporary password and group files
link_passwd();


######## Install as a distribution package ########
# install packages in directory to tar
# dereference array PKGS references
if ($config->{PKGS}) {
    my @packages = @{$config->{PKGS}};
    for my $p (@packages) {
	print "$pkg install $p\n";
	system("$pkg $p\n");
    }
}

# dereference hash CUSTOM references
if ($config->{CUSTOM}) {
    my %customs = %{$config->{CUSTOM}};
    for my $c (keys %customs) {
	my $val = %customs{$c};

	# make the directory to install to
	my $valdir = `dirname $pkgroot/$val`;
	system("mkdir -p $valdir");

	print "cp -r $c $pkgroot/$val\n";
	system("cp -r $c $pkgroot/$val\n");
    }
}

######## Install on the Live CD ########
if ($config->{LIVE_CD_PKGS}) {
    my @packages = @{$config->{LIVE_CD_PKGS}};
    for my $p (@packages) {
	print "$pkg install $p\n";
	system("pkg -r $workroot install -y $p\n");
    }
}

if ($config->{LIVE_CD_CUSTOM}) {
    my %customs = %{$config->{LIVE_CD_CUSTOM}};
    for my $c (keys %customs) {
	my $val = %customs{$c};

	# make the directory to install to
	my $valdir = `dirname $workroot/$val`;
	system("mkdir -p $valdir");

	print "cp -r $c $workroot/$val\n";
	system("cp -r $c $workroot/$val\n");
    }
}

my $installerconfig = $config->{INSTALLERCONFIG};
if ($installerconfig) {
    system("cp $installerconfig $workroot/etc/");
}

# pack up the custom stuff as a distribution tar
print "--------------------------------------------------\n";
print "Creating distribution file $distdir/$distname...\n";
print "--------------------------------------------------\n";

my $startdir = getcwd;

chdir "$pkgroot" || die "cannot change to $pkgroot\n";
system("tar -cvJpf $distdir/$distname ./");
chdir "$startdir" || die "cannot change back to start dir\n";

# add $distname to MANIFEST
manifest_add;

# check if DEBUG is defined
if ($DEBUG) {
    # stop so we can check out the mounted files
    print "--------------------------------------------------\n";
    print "DEBUG flag set. Not generating iso image.\n";
    print "--------------------------------------------------\n";
    exit(0);
}

# pretty self explanatory by the function name
print "--------------------------------------------------\n";
write_new_img;
print "--------------------------------------------------\n";

# unmount and free vnode
print "--------------------------------------------------\n";
print "Deleting $workroot...\n";
print "--------------------------------------------------\n";
remove_dirs;
unmount_img;

# bye
exit(0);
