#!/usr/bin/perl

# draklive $Id$

# Copyright (C) 2005 Mandriva
#                    Olivier Blin <oblin@mandriva.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2, or (at your option)
# any later version.
#
# 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; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

use lib qw(/usr/lib/libDrakX);
use MDK::Common;
use list_modules;
use modules;
use run_program;
use POSIX qw(strftime);
use Cwd 'abs_path';
use Getopt::Long;
use Pod::Usage;

my $dir_distrib_sqfs = {
    mountpoint => '/distrib',
    type => 'squashfs',
    source => 'distrib.sqfs',
    build_from => '/',
};
my $dir_memory = {
    mountpoint => '/memory',
    type => 'tmpfs',
};

my %predefined = (
    mounts => {
        simple_union => {
            root => '/union',
            overlay => 'unionfs',
            dirs => [
                $dir_memory,
                { mountpoint => '/media' },
            ],
        },
        squash_union => {
            root => '/union',
            overlay => 'unionfs',
            dirs => [
                {
                    mountpoint => '/system',
                    type => 'loopfs',
                    pre_allocate => '100k',
                    source => 'system.loop'
                },
                $dir_distrib_sqfs,
            ],
        },
        volatile_squash_union => {
            root => '/union',
            overlay => 'unionfs',
            dirs => [
                $dir_memory,
                $dir_distrib_sqfs,
            ],
        },
        multi_squash_union => {
            root => '/union',
            overlay => 'unionfs',
            dirs => [
                $dir_memory,
                {
                    mountpoint => '/system',
                    type => 'squashfs',
                    source => 'system.sqfs'
                },
                $dir_distrib_sqfs,
            ],
        },
    },
    media => {
        usb => {
            storage => 'usb',
            fs => 'vfat',
            sleep => 15,
            source => 'LABEL=MDVUSBROOT',
            mountpoint => '/media',
        },
        cdrom => {
            storage => 'cdrom',
            fs => 'iso9660',
            source => 'LABEL=MDVCDROOT',
            mountpoint => '/media',
        },
    },
);

my %custom = (
    media => {
        nfs => sub {
            my ($module, $client, $source) = @_;
            {
                modules => [ $module ],
                fs => 'nfs',
                pre => "ifconfig eth0 $client up",
                source => $source,
                mountpoint => '/media',
            };
        },
    },
);

sub nls_modules {
    my ($live) = @_;
    if_($live->{media}{fs} eq 'vfat', 'nls_cp437'), #- default FAT codepage
    map { "nls_$_" } (map { "iso8859-$_" } 1..7, 9, 13..15), 'utf8';
}

sub fs_module {
    my ($live) = @_;
    $live->{media}{fs} eq 'iso9660' ? 'isofs' : $live->{media}{fs};
}

sub progress_start {
    my ($total, $time, $o_exp_divide) = @_;
    {
        total => $total,
        current => 0,
        start_time => $time,
        exp_divide => $o_exp_divide,
        maxl => length($total) - $o_exp_divide,
    };
}

sub progress_show_incr {
    my ($progress, $incr, $time) = @_;
    $progress->{current} += $incr;
    my $elapsed_time = $time - $progress->{start_time};
    my $eta = int($elapsed_time*$progress->{total}/$progress->{current});
    printf("\r%3d%% (%$progress->{maxl}s/%-$progress->{maxl}s), %8s/%8s (ETA)",
           int(100*$progress->{current}/$progress->{total}),
           (map { substr($_, 0, length($_)-$progress->{exp_divide}) } $progress->{current}, $progress->{total}),
           (map { strftime("%H:%M:%S", gmtime($_)) } $elapsed_time, $eta));
}

sub progress_end() { print "\n" }

my $loop_number = 0;
my %loop = (
    squashfs => {
        read_only => 1,
        modules => [ qw(loop squashfs) ],
        build => sub {
            my ($dest, $root) = @_;
            my $total = first(split /\s/, `du -sb $root`);
            print "have to process " . int($total/1000000) . " MB\n";
            my $progress = progress_start($total, time(), 6);
            open(my $OUTPUT, '-|', 'mksquashfs', $root, $dest, '-info');
            {
                local $_; #- avoid outside $_ to be overwritten
                while (<$OUTPUT>) {
                    if (/^mksquashfs: file .*, uncompressed size (\d+) bytes, (?:DUPLICATE)?$/) {
                        progress_show_incr($progress, $1, time());
                    }
                }
            }
            progress_end();
        },
        mount => sub {
            my ($live, $dir) = @_;
            my @mnt = (
                ($live->{media}{fs} eq 'nfs' ? '/bin/losetup -r' : 'losetup') .
                  " /dev/loop$loop_number $live->{media}{mountpoint}/$dir->{source}",
                "mount -o ro -t squashfs /dev/loop$loop_number $dir->{mountpoint}");
            $loop_number++;
            @mnt;
        },
    },
    loopfs => {
        modules => [],
        create => sub {
            my ($dest, $size) = @_;
            run_('dd', "of=$dest", 'count=0', "seek=$size", 'bs=1k');
            run_('mke2fs', '-F', $dest);
        },
        mount => sub {
            my ($live, $dir) = @_;
            my @mnt = (
                "losetup /dev/loop$loop_number $live->{media}{mountpoint}/$dir->{source}",
                "mount -t ext2 /dev/loop$loop_number $dir->{mountpoint}");
            $loop_number++;
            @mnt;
        },
    },
    tmpfs => {
        mount => sub {
            my ($_live, $dir) = @_;
            "mount -t tmpfs none $dir->{mountpoint}";
        },
    },
);

my %overlay = (
    unionfs => {
        modules => [ qw(unionfs) ],
        mount => sub {
            my ($live) = @_;
            #- build dirs list: "dir1=ro:dir2:ro:dir3=rw"
            my $dirs = join(':',
                            map { "$_->{mountpoint}=" .
                                    ($_->{type} && !$loop{$_->{type}}{read_only} ? 'rw' : 'ro');
                                } @{$live->{mount}{dirs} || []});
            "mount -o dirs=$dirs -t unionfs unionfs $live->{mount}{root}";
        },
    },
);

my %storage = (
    cdrom => {
        read_only => 1,
        modules => 'disk/cdrom|hardware_raid|sata|scsi bus/usb disk/raw',
    },
    usb => {
        modules => 'bus/usb disk/raw|usb',
        create => undef,
        record => \&record_usb_master,
    },
);

my %moddeps;
sub load_moddeps {
    my ($root, $kernel_path) = @_;
    my $get_modname = sub { first($_[0] =~ m!^$kernel_path/kernel/(?:.*/|)(.*?)\.k?o!) };
    %moddeps = (map {
	my ($f, $deps) = split ':';
	my $modname = $get_modname->($f);
	$modname => { full => $f, deps => [ map { $get_modname->($_) } split ' ', $deps ] };
    } cat_($root . $kernel_path . '/modules.dep'));
}
sub moddeps_closure {
    my ($module) = @_;
    my @deps = @{$moddeps{$module}{deps}};
    (map { moddeps_closure($_) } @deps), @deps;
}

sub run_ {
    print "running " . join(' ', @_) . "\n";
    run_program::run(@_);
}

sub create_initrd {
    my ($live) = @_;

    rm_rf($live->{initrd_tree}) if -e $live->{initrd_tree};

    mkdir_p($live->{initrd_tree} . $_) foreach
      qw(/bin /dev /lib /proc /sys),
      $live->{media}{mountpoint},
      (map { $_->{mountpoint} } @{$live->{mount}{dirs} || []}),
      $live->{mount}{root};

    # cp_f($live->{system}{root} . '/sbin/nash', $live->{initrd_tree} . '/bin/');
    #- use nash from cooker for now, label support
    cp_f('/sbin/nash', $live->{initrd_tree} . '/bin/');

    if ($live->{debug} || $live->{media}{fs} eq 'nfs') {
        cp_f('/lib/tls/libc.so.6', $live->{initrd_tree} . '/lib/');
        cp_f('/lib/ld-linux.so.2', $live->{initrd_tree} . '/lib/');
    }
    if ($live->{media}{fs} eq 'nfs') {
        cp_f('/sbin/ifconfig', $live->{initrd_tree} . '/bin/');
        cp_f('/bin/mount', $live->{initrd_tree} . '/bin/');
        cp_f('/sbin/losetup', $live->{initrd_tree} . '/bin/');
        if ($live->{debug}) {
            cp_f('/bin/ping', $live->{initrd_tree} . '/bin/');
            cp_f('/lib/libresolv.so.2', $live->{initrd_tree} . '/lib/');
        }
    }
    if ($live->{debug}) {
        cp_f('/usr/bin/strace', $live->{initrd_tree} . '/bin/');
        cp_f('/usr/bin/busybox', $live->{initrd_tree} . '/bin');
        my @l = map { /functions:/ .. /^$/ ? do { s/\s//g; split /,/ } : () } `busybox`;
        shift @l;
        symlink('busybox', $live->{initrd_tree} . "/bin/$_") foreach @l;
    }

    require devices;
    devices::make($live->{initrd_tree} . "/dev/$_") foreach
        qw(console initrd null ram systty),
        (map { "tty$_" } 0..5),
        (map { "loop$_" } 0..7),
        (map { $_, $_ . '1' } map { "sd$_" } ('a' .. 'h'));

    load_moddeps($live->{system}{root}, "/lib/modules/" . $live->{system}{kernel});

    my ($modules, $skipped) = partition { exists $moddeps{$_} }
      uniq(map { modules::cond_mapping_24_26($_) } category2modules($storage{$live->{media}{storage}}{modules}));
    my ($extra_modules, $missing) = partition { exists $moddeps{$_} }
      nls_modules($live),
      fs_module($live),
      @{$live->{media}{modules} || []},
      (map { @{$loop{$_}{modules} || []} } uniq(map { $_->{type} } grep { $_->{type} } @{$live->{mount}{dirs} || []})),
      ($live->{mount}{overlay} ? @{$overlay{$live->{mount}{overlay}}{modules} || []} : ());

    @$missing and die "missing mandatory modules:" . join("\n", '', sort(@$missing));
    push @$modules, @$extra_modules;

    my @module_deps = uniq(map { moddeps_closure($_) } @$modules);
    run_('gzip', '>', $live->{initrd_tree} . "/lib/$_.ko", '-dc', $live->{system}{root} . $moddeps{$_}{full})
      foreach @module_deps, @$modules;

    @$skipped and warn "skipped modules:" . join("\n", '', sort(@$skipped));

    create_initrd_linuxrc($live, @module_deps, @$modules);
    compress_initrd_tree($live);
    add_splash($live);
    $live->{copy_initrd} and cp_f($live->{boot_dir} . '/initrd.gz', $live->{copy_initrd});
}

sub create_initrd_linuxrc {
    my ($live, @modules) = @_;
    my $target = $live->{mount}{root} || $live->{media}{mountpoint};
    output_with_perm($live->{initrd_tree} . '/linuxrc', 0755,
                     join("\n",
                          "#!/bin/nash",
                          (map { "insmod /lib/$_.ko" } @modules),
                          if_($live->{media}{sleep}, "sleep $live->{media}{sleep}"),
                          #- required for labels
                          "mount -t proc none /proc",
                          #- required for cdrom labels
                          "mount -t sysfs none /sys",
                          if_($live->{debug}, "/bin/sh"),
                          if_($live->{media}{pre}, deref_array($live->{media}{pre})),
                          ($live->{media}{fs} eq 'nfs' ? '/bin/mount -n -o ro,nolock' : 'mount') .
                            ($storage{$live->{media}{storage}}{read_only} && " -o ro") .
                            " -t $live->{media}{fs} $live->{media}{source} $live->{media}{mountpoint}",
                          (map { $loop{$_->{type}}{mount}->($live, $_) } grep { $_->{type} } @{$live->{mount}{dirs} || []}),
                          ($live->{mount}{overlay} ? $overlay{$live->{mount}{overlay}}{mount}->($live) : ()),
                          "echo 0x0100 > /proc/sys/kernel/real-root-dev",
                          "umount /sys",
                          "umount /proc",
                          ($live->{mount}{overlay} ?
                             # don't move to /initrd but /live, or else the overlay will be unmounted
                             ("mkdir -p $target/live", "pivot_root $target $target/live") :
                               "pivot_root $target $target/initrd"),
                          if_($live->{post}, deref_array($live->{post})),
                          ""));
}

sub compress_initrd_tree {
    my ($live) = @_;

    my $size = run_program::get_stdout("du -ks $live->{initrd_tree} | awk '{print \$1}'") + 250;
    my $inodes = run_program::get_stdout("find $live->{initrd_tree} | wc -l") + 1250;
    $size = int($size + $inodes / 10) + 1; #- 10 inodes needs 1K
    my $initrd = $live->{boot_dir} . "/initrd";

    mkdir_p($live->{boot_dir});
    run_('dd', 'if=/dev/zero', "of=$initrd", 'bs=1k', "count=$size");
    run_('mke2fs', '-q', '-m', 0, '-F', '-N', $inodes, '-s', 1, $initrd);
    mkdir_p($live->{mnt});
    run_('mount', '-o', 'loop', '-t', 'ext2', $initrd, $live->{mnt});
    cp_af(glob("$live->{initrd_tree}/*"), $live->{mnt});
    rm_rf($live->{mnt} . "/lost+found");
    run_('umount', $live->{mnt});
    run_('gzip', '-f', '-9', $initrd);
}

sub add_splash {
    my ($live) = @_;
    if ($live->{system}{vga_mode} && $live->{system}{splash} ne 'no') {
	require bootloader;
	my $initrd = "$live->{boot_dir}/initrd.gz";
	my $tmp_initrd = '/tmp/initrd.gz';
        cp_f($initrd, $live->{system}{root} . $tmp_initrd);
        {
            local $::prefix = $live->{system}{root};
            bootloader::add_boot_splash($tmp_initrd, $live->{system}{vga_mode});
        }
        cp_f($live->{system}{root} . $tmp_initrd, $initrd);
	unlink($live->{system}{root} . $tmp_initrd);
    }
}

sub build_syslinux_cfg {
    my ($live) = @_;
    #- fastboot is needed to avoid fsck
    my $append = "fastboot splash=silent vga=$live->{system}{vga_mode}";
    qq(default live
prompt 1
timeout 40
display live.msg
label live
    kernel vmlinuz
    append initrd=initrd.gz $append
);
}

sub install_system {
    my ($live) = @_;
    run_('drakx-in-chroot',
         $live->{system}{repository},
         $live->{system}{root},
         if_($live->{system}{auto_install}, '--auto_install', abs_path($live->{system}{auto_install})),
         if_($live->{system}{patch}, '--defcfg', abs_path($live->{system}{patch})))
      or die "unable to install system chroot";

    run_('urpmi',
         '--root',
         $live->{system}{root},
         map { abs_path($_) } @{$live->{system}{rpms}}) if @{$live->{system}{rpms}};

    #- make sure harddrake is run
    #- (do it in chroot, or else Storable from the build box may write an incompatible config file)
    system("chroot $live->{system}{root} " .
           "perl -MStorable -e \"Storable::store({ UNKNOWN => {} }, '/etc/sysconfig/harddrake2/previous_hw')\"");

    #- interactive mode can lead to race in initscripts
    #- (don't use addVarsInSh from MDK::Common, it breaks shell escapes)
    substInFile { s/^PROMPT=.*/PROMPT=no/ } $live->{system}{root} . '/etc/sysconfig/init';

    #- disable first boot wizard
    output($live->{system}{root} . '/etc/sysconfig/firstboot', 'FIRSTBOOT=no');
    #- enable drakx-finish-install
    output($live->{system}{root} . '/etc/sysconfig/finish-install', 'FINISH_INSTALL=yes');

    #- preselect guest user in kdm
    my $kdm_cfg = '/usr/share/config/kdm/kdmrc';
    update_gnomekderc($live->{system}{root} . $kdm_cfg,
                      'X-:0-Greeter' => (PreselectUser => 'Default', DefaultUser => 'guest')) if -f $kdm_cfg;
}

sub create_loopback_files {
    my ($live) = @_;
    mkdir_p($live->{images_dir});
    foreach (grep { $_->{build_from} } @{$live->{mount}{dirs} || []}) {
        my $tree = $live->{system}{root} . $_->{build_from};
        my $dest = $live->{images_dir} . '/' . $_->{source};
        unlink($dest);
        $loop{$_->{type}}{build}->($dest, $tree);
    }
    foreach (grep { $_->{pre_allocate} } @{$live->{mount}{dirs} || []}) {
        my $dest = $live->{images_dir} . '/' . $_->{source};
        unlink($dest);
        $loop{$_->{type}}{create}->($dest, $_->{pre_allocate});
    }
}

sub get_media_label {
    my ($live) = @_;
    first($live->{media}{source} =~ /^LABEL=(.*)$/);
}

sub get_media_device {
    my ($live) = @_;
    return $live->{media}{device} if $live->{media}{device};
    my $label = get_media_label($live) or return $live->{media}{source};
    my $device = chomp_(`readlink -f /dev/disk/by-label/$label`)
      or die "unable to find device for /dev/disk/by-label/$label";
    $device;
}

sub prepare_bootloader {
    my ($live) = @_;
    mkdir_p($live->{boot_dir});
    cp_f($live->{system}{root} . '/boot/vmlinuz-' . $live->{system}{kernel}, $live->{boot_dir} . '/vmlinuz');
    my $msg = $live->{system}{root} . '/boot/message-graphic';
    cp_f($msg, $live->{boot_dir} . '/live.msg') if -f $msg;
    if ($live->{media}{storage} eq 'usb') {
        output($live->{boot_dir} . '/syslinux.cfg', build_syslinux_cfg($live));
    } elsif ($live->{media}{storage} eq 'cdrom') {
        cp_f('/usr/lib/syslinux/isolinux-graphic.bin', $live->{boot_dir} . '/isolinux.bin');
        output($live->{boot_dir} . '/isolinux.cfg', build_syslinux_cfg($live));
    } else {
        warn "not implemented yet";
    }
}

sub create_master {
    my ($live) = @_;
    if (my $create = $storage{$live->{media}{storage}}{create}) {
        $create->($live, $o_refresh_boot_only);
    } else {
        warn "not implemented yet";
    }
}

sub record_usb_master {
    my ($live, $o_refresh_boot_only) = @_;
    my $label = get_media_label($live);
    if ($live->{media}{device} && $label) {
        run_('mkdosfs', '-n', $label, $live->{media}{device})
          or die "unable to format device $live->{media}{device}";
    }
    my $device = get_media_device($live);
    mkdir_p($live->{mnt});
    run_('mount', $device, $live->{mnt})
      or die "unable to mount $device";
    cp_f(map { $live->{boot_dir} . "/$_" } qw(vmlinuz initrd.gz live.msg), $live->{mnt});
    unless ($o_refresh_boot_only) {
        foreach (grep { $_->{build_from} || $_->{pre_allocate} } @{$live->{mount}{dirs} || []}) {
            print "copying $_->{source}\n";
            cp_f($live->{images_dir} . '/' . $_->{source}, $live->{mnt});
        }
    }
    run_('umount', $live->{mnt});
    #- use syslinux -s, "safe, slow and stupid" version of SYSLINUX
    run_('syslinux', '-s', $device) or die "unable to run syslinux on $device";
}

sub record_master {
    my ($live, $o_refresh_boot_only) = @_;
    if (my $record = $storage{$live->{media}{storage}}{record}) {
        $record->($live, $o_refresh_boot_only);
    } else {
        warn "not implemented yet";
    }
}

sub record_boot {
    my ($live) = @_;
    record_master($live, 'refresh');
}

sub complete_config {
    my ($live) = @_;

    #- set unsupplied config dirs
    $live->{workdir} ||= '/tmp/draklive';
    $live->{boot_dir} ||= $live->{workdir} . "/boot";
    $live->{initrd_tree} ||= $live->{workdir} . "/initrd";
    $live->{images_dir} ||= $live->{workdir} . "/images";
    $live->{mnt} ||= $live->{workdir} . "/mnt";

    #- check for minimum requirements
    ref $live->{media} or die "no media definition";
    ref $live->{system} or die "no system definition";
    $live->{system}{kernel} or die "no kernel has been configured";
    mkdir_p($live->{workdir});
}

sub clean {
    my ($live) = @_;
    rm_rf($_) foreach grep { -e $_ } $live->{workdir}, $live->{system}{root};
}

my @actions = (
    { name => 'clean', do => \&clean },
    { name => 'install', do => \&install_system },
    { name => 'initrd', do => \&create_initrd },
    { name => 'boot', do => \&prepare_bootloader },
    { name => 'loop', do => \&create_loopback_files },
    { name => 'master', do => \&create_master },
    { name => 'record', do => \&record_master },
    { name => 'record_boot', do => \&record_boot },
);
my @all = qw(install initrd boot loop master);

my %live;
GetOptions(
    "help" => sub { pod2usage('-verbose' => 1) },
    "all" => sub { $_->{to_run} = 1 foreach grep { member($_->{name}, @all) } @actions },
    (map { $_->{name} => \$_->{to_run} } @actions),
    "device:s" => sub { $live{media}{device} = $_[1] },
    "config:s" => sub {
        my $path = $_[1];
        #- don't use do(), since it can't see lexicals in the enclosing scope
        my $cfg = eval(cat_($path)) or die "unable to load $path";
        add2hash(\%live, $cfg);
        print "loaded $path as config file\n";
    },
) or pod2usage();
unless (keys(%live)) {
    warn 'no live definition';
    pod2usage();
}
complete_config(\%live);

require standalone;
every { !$_->{to_run} } @actions and die 'nothing to do';
foreach (grep { $_->{to_run} } @actions) {
    print qq(* entering step "$_->{name}"\n);
    $_->{do}->(\%live);
    print qq(* step "$_->{name}" done\n);
}

__END__

=head1 NAME

draklive - A live distribution mastering tool

=head1 SYNOPSIS

draklive [options]

 Options:
   --help           long help message

   --install        install selected distribution in chroot
   --initrd         build initrd
   --boot           prepare bootloader files
   --loop           build compressed loopback files
   --master         build master image

   --all            run all steps, from installation to mastering

   --clean          clean installation chroot and work directory
   --record         install live on selected media
   --record_boot    install bootloader only on selected media

   --device <dev>   use this device for live recording, formatting
                    it preliminary (not needed if the device
                    already has the required label)
   --config <file>  use this configuration file as live description

Examples:

 draklive --config config/live.cfg --clean

 draklive --config config/live.cfg --all

 draklive --config config/live.cfg --record --device /dev/sdb1

=head1 OPTIONS

=over 8

=item B<--config>

Makes draklive use the next argument as a configuration file.
This file should contain an hash describing the live distribution,
meaning the system (chroot and boot), media (usb, cdrom, nfs),
and mount type (simple R/W union, union with squash files).

Here's a configuration sample:

  {
    system => {
        root => '/chroot/live-move',
        repository => '/mnt/ken/2006.0/i586',
        kernel => '2.6.12-12mdk-i586-up-1GB',
        auto_install => 'config/auto_inst.cfg.pl',
        patch => 'config/patch-2006-live.pl',
        rpms => [
             'rpms/unionfs-kernel-2.6.12-12mdk-i586-up-1GB-1.1.1.1.20051124.1mdk-1mdk.i586.rpm'
        ],
        vga_mode => 788,
    },
    media => $predefined{media}{usb},
    mount => $predefined{mounts}{squash_union}
  };

=back

=head1 DESCRIPTION

B<draklive> builds a live distribution according to a
configuration file, creates a master image,
and optionnally installs it on a device.

See L<http://qa.mandriva.com/twiki/bin/view/Main/DrakLive>

=head1 AUTHOR

Olivier Blin <oblin@mandriva.com>

=cut