#!/usr/bin/perl -w

# radmindWrapper
# ==============
# Ian Ward Comfort, 25 October 2006
#
# This script Radminds a machine back to its pristine state. It executes all
# the necessary Radmind tools (ktcheck, fsdiff, lapply) and runs a few post-
# flight actions. It may be run manually (with sudo) or installed in root's
# crontab (as it is for most of our managed machines).
#
# When run with no options, the machine will be Radminded immediately (using
# iHook to lock out the console and provide feedback) and then rebooted. This
# makes it easy for a tech to simply type 'sudo radmindWrapper' on a cluster
# machine and walk away, or for a remote administrator to start a Radmind
# session with a one-line command sent via SSH.
#
# When run daily by cron, we use the --stagger option so that our machines
# don't hit ResXserve all at once. We also use the --atlogout flag to wait
# until no one is logged into the machine to Radmind. (The script itself won't
# actually wait; it sets a trigger file which tells our logout hook to run the
# script again when the user logs out.)
#
# Note that radmindWrapper will attempt to update itself before it updates the
# rest of the filesystem. radmindWrapper also does its best to ensure a
# successful update even when critical system files (including the kernel and
# associated libraries) require replacement. It attempts to copy the libraries
# and frameworks required by the tools it uses into a temporary directory if
# they are being replaced by lapply. (Of course any major updates should be
# thoroughly tested before they are applied to production machines.) Many
# thanks to Andrew Mortensen of the University of Michigan for the tool-saving
# strategy which is automated here. I've added the chroot logic to deal with
# updates that replace /usr/lib/dyld with an incompatible version.

use Getopt::Long;
use File::Temp qw(tempdir);
use File::Basename;
use POSIX qw(strftime setsid);
use Sys::Hostname;
use Sys::Syslog;

use strict;
use warnings;

# Constants and initial setup

my $SAFE_TMPDIR            = tempdir('/tmp/radmindWrapper.XXXXXX');
my $TOOL_DIR               = "$SAFE_TMPDIR/tools";

$ENV{'PATH'} = "$TOOL_DIR:/usr/local/bin:/usr/local/sbin:/bin:/sbin:/usr/bin:/usr/sbin";
$ENV{'DYLD_LIBRARY_PATH'} = $ENV{'DYLD_FRAMEWORK_PATH'} = $TOOL_DIR;

my @TOOLS_NEEDED           = ('/bin/sh', '/sbin/reboot', '/usr/bin/touch', '/usr/local/bin/PlistBuddy');
my $DYLD_PATH              = '/usr/lib/dyld';
my $DYLD_BEING_REPLACED    = 0; # Will be toggled if dyld appears in an applicable transcript

my @PRE_UPDATE_DIRS        = ('/usr/local/bin', '/Library/Management');
my $WRAPPER_PATH           = '/usr/local/bin/radmindWrapper';
my $WRAPPER_BEING_REPLACED = 0; # Will be toggled if radmindWrapper appears in an applicable transcript

my $RADMIND_SERVER         = "resxserve.stanford.edu";
my $RADMIND_PORT           = 6222;
my $APPLICABLE_TRANSCRIPT  = "$SAFE_TMPDIR/applicable.T";
my $COMPARISON_PATH        = "."; # Should be either . or /

my $KTCHECK_FLAGS          = "-C -v -c sha1 -h $RADMIND_SERVER -p $RADMIND_PORT";
my $FSDIFF_FLAGS           = "-A -c sha1 -o $APPLICABLE_TRANSCRIPT -%";
my $LAPPLY_FLAGS           = "-F -i -% -c sha1 -h $RADMIND_SERVER -p $RADMIND_PORT $APPLICABLE_TRANSCRIPT";

my $MAX_STAGGER_MINUTES    = 150;
my $LOGOUT_TRIGGER_FILE    = "/Library/Management/.RunRadmindAtLogout";
my $IHOOK_LAUNCHD_PLIST    = "/Library/Management/edu.stanford.rescomp.radlock.plist";

my $PID_FILE               = "/var/run/radmindWrapper.pid";
my $LOG_FILE               = "/var/log/radmindWrapper.log";
my $ERROR_FILE             = "/var/log/radmindWrapper.error";
my $LOG_FACILITY           = "user";
local $\                   = "\n";

# Some functions we'll need

sub writepid {
    open PID, ">$PID_FILE" or warn "Can't write to $PID_FILE: $!";
    print PID $$;
    close PID;
}

sub daemonize {
    chdir '/'                 or die "Can't chdir to /: $!";
    open STDIN, '/dev/null'   or die "Can't read /dev/null: $!";
    open STDOUT, '>/dev/null' or die "Can't write to /dev/null: $!";
    defined (my $pid = fork)  or die "Can't fork: $!";
    exit if $pid;
    setsid                    or die "Can't start a new session: $!";
    open STDERR, '>&STDOUT'   or die "Can't dupe stdout: $!";
    writepid;
}

sub uniq {
    my %seen;
    grep { !$seen{$_}++ } @_;
}

sub system_safe {
    if ($DYLD_BEING_REPLACED) {
        defined (my $pid = fork) or die "Can't fork: $!";
        if ($pid) {
            waitpid($pid, 0);
            return $?;
        } else {
            $SIG{INT} = $SIG{TERM} = $SIG{HUP} = $SIG{__DIE__} = $SIG{__WARN__};
            chroot $TOOL_DIR or die "Can't chroot to $TOOL_DIR: $!";
            $ENV{'PATH'} = $ENV{'DYLD_LIBRARY_PATH'} = $ENV{'DYLD_FRAMEWORK_PATH'} = '/';
            system @_;
            exit;
        }
    } else {
        system @_;
    }
}

sub cleanup {
    syslog('err', '%s', $_[0]) if (defined $_[0]);
    closelog;
    system("launchctl", "unload", $IHOOK_LAUNCHD_PLIST);
    unlink $PID_FILE;
}

# Defaults for command line options
my $staggered_start = 0;
my $run_at_logout   = 0;
my $reboot_after    = 1;
my $use_ihook       = 1;
my $stay_attached   = 0;
my $print_usage     = 0;

Getopt::Long::Configure('bundling');
GetOptions('stagger|s!'       => \$staggered_start,
           'atlogout'         => \$run_at_logout,
           'reboot|r!'        => \$reboot_after,
           'ihook|i!'         => \$use_ihook,
           'verbose|v!'       => \$stay_attached,
           'usage|u|help|h|?' => \$print_usage);

if ($print_usage) {
    print "radmindWrapper: use the Radmind tools to restore a machine to spec";
    print "    Options (most can be negated with a prefix of 'no'):";
    print "      --stagger, -s     wait between 0 and $MAX_STAGGER_MINUTES minutes to start";
    print "      --atlogout        wait until no user is at the console";
    print "      --reboot, -r      reboot this machine after Radminding (on by default)";
    print "      --ihook, -i       use iHook to lock the console while running (on by default)";
    print "      --verbose, -v     don't dissociate from terminal";
    print "      --help, -h, -?    print this message";
    print "    No options is equivalent to using '--nostagger --reboot --ihook'.";
    exit;
}

die "This tool must be run as root. Try 'radmindWrapper --help' for options.\n" if ($> != 0);

openlog('radmindWrapper', 'pid', $LOG_FACILITY);
$SIG{__WARN__} = sub { syslog('warning', '%s', $_[0]) };

if (-f $PID_FILE) {
    open PID, $PID_FILE or warn "Can't read $PID_FILE: $!";
    chomp(my $pid = <PID>);
    close PID;
    die "radmindWrapper is already running with pid $pid.\n" if (kill 0 => $pid);
}

# Set __DIE__ handler after the test above, so we don't cleanup another instance's iHook job
$SIG{INT} = $SIG{TERM} = $SIG{HUP} = $SIG{__DIE__} = \&cleanup;

writepid;

chdir '/' or die "Can't chdir to /: $!";

# Disable sleep while Radminding or waiting to Radmind
system "pmset sleep 0";

if ($staggered_start) {
    my $minutes_to_wait = rand $MAX_STAGGER_MINUTES;
    printf "Staggered start time was requested, waiting for about %d minutes...\n", (int $minutes_to_wait);
    daemonize unless ($stay_attached);
    sleep 60 * $minutes_to_wait;
}

if ($run_at_logout) {
    open WHO, "who |" or die "Can't read from who: $!";
    my $users = do { local $/; <WHO> };
    close WHO;
    if ($users =~ / console /) {
        print "User is at the console, setting trigger file and exiting.";
        system("touch", $LOGOUT_TRIGGER_FILE);
        cleanup;
        exit;
    }
}

if ($stay_attached) {
    open STDOUT, "| tee $LOG_FILE" or warn "Can't tee to $LOG_FILE: $!";
} else {
    daemonize unless ($staggered_start);
    open STDOUT, ">$LOG_FILE" or warn "Can't write to $LOG_FILE: $!";
}

open STDERR, ">$ERROR_FILE" or warn "Can't write to $ERROR_FILE: $!";
$| = 1;

if ($use_ihook) {
    # Keep ktcheck output from going to iHook
    $KTCHECK_FLAGS .= " 1>&2";
    (system("launchctl", "load", $IHOOK_LAUNCHD_PLIST) == 0) or die "Can't load job $IHOOK_LAUNCHD_PLIST: $?";
}

print "%BEGINPOLE" if ($use_ihook);
print "Updating command files and transcripts...";
my $return_value = system("ktcheck $KTCHECK_FLAGS") >> 8;
die "ktcheck returned $return_value" unless ($return_value <= 1);

# We seem to have the right hostname, so copy it to other fields
(my $hostname = hostname) =~ s/\..*//;
system("scutil", "--set", "ComputerName", $hostname);
system("scutil", "--set", "LocalHostName", $hostname);

# Try to update our imaging tools (including this script) first
print "Updating imaging tools...";
my $prefix = ($COMPARISON_PATH eq "/" ? "" : ".");

foreach my $comp_path (@PRE_UPDATE_DIRS) {
    $return_value = system("fsdiff $FSDIFF_FLAGS $prefix$comp_path") >> 8;
    $return_value = system("lapply $LAPPLY_FLAGS") >> 8 unless ($return_value);
    unless ($return_value) {
        open TRANSCRIPT, $APPLICABLE_TRANSCRIPT;
        my $transcript = do { local $/; <TRANSCRIPT> };
        close TRANSCRIPT;
        $WRAPPER_BEING_REPLACED = 1 if ($transcript =~ m/^(\+ .|h|l) \Q$prefix$WRAPPER_PATH/m);
    }
}

if ($WRAPPER_BEING_REPLACED) {
    my $wrapper_command = $WRAPPER_PATH;
    $wrapper_command .= ' -i' if ($use_ihook);
    $wrapper_command .= ' -r' if ($reboot_after);
    $wrapper_command .= ' -v' if ($stay_attached);
    { unlink $PID_FILE and exec $wrapper_command; }
    warn "Couldn't launch new wrapper script, continuing anyway";
}

print "Examining system for necessary updates...";
$return_value = system("fsdiff $FSDIFF_FLAGS $COMPARISON_PATH") >> 8;
die "fsdiff returned $return_value" unless ($return_value == 0);

print "%BEGINPOLE" if ($use_ihook);
print "Preparing to update files...";

# Save current datestamp for restoration if fsdiff fails
chomp(my $last_datestamp = `PlistBuddy -c 'print :LoginwindowText' /Library/Preferences/com.apple.loginwindow.plist`);

my $clear_ls_cache         = 0;
my $clear_components_cache = 0;
my $rebuild_kernel_caches  = 0;
my (%seen, %links_to_create);

print STDERR strftime('%+', localtime) . " Calculating tool dependencies";

# Generate list of dependencies for tools needed after lapply
foreach my $binary (@TOOLS_NEEDED) {
    # Add target if symlink
    if (-l $binary and defined(my $linkdest = readlink $binary)) {
        (ord $linkdest == ord '/') or $linkdest = dirname($binary) . '/' . $linkdest;
        push @TOOLS_NEEDED, $linkdest unless $seen{$linkdest}++;
        $links_to_create{basename $binary} = basename $linkdest;
    }

    # Add dynamically loaded libraries
    open OTOOL, "-|", "otool", "-L", $binary or die "Can't read from otool: $!";
    <OTOOL> or die "Can't read from otool: $!";
    while (<OTOOL>) {
        my ($garbage, $dependency) = split /\s+/;
        push @TOOLS_NEEDED, $dependency unless $seen{$dependency}++;
    }
    close OTOOL;
}

map { s/(?<=\.framework).*// } @TOOLS_NEEDED;
@TOOLS_NEEDED = uniq @TOOLS_NEEDED;

print STDERR strftime('%+', localtime) . " Finished calculating, tools needed:";
print STDERR "  $_" foreach @TOOLS_NEEDED;
print STDERR strftime('%+', localtime) . " Copying tools that are being replaced";

# Examine transcript to see if we need to save tools or clear caches
open TRANSCRIPT, $APPLICABLE_TRANSCRIPT or die "Can't read applicable transcript: $!";
my $transcript = do { local $/; <TRANSCRIPT> };
close TRANSCRIPT;

study $transcript;

$clear_ls_cache         = 1 if $transcript =~ m{com\.apple\.LaunchServices};
$clear_components_cache = 1 if $transcript =~ m{ \Q$prefix\E/System/Library/(Components|QuickTime)};
$rebuild_kernel_caches  = 1 if $transcript =~ m{ \Q$prefix\E/System/Library/Extensions};

if ($transcript =~ m/^(. .|h|l) \Q$prefix$DYLD_PATH/m) {
    (system("ditto", $DYLD_PATH, "$TOOL_DIR$DYLD_PATH") == 0) or die "Can't ditto $DYLD_PATH to $TOOL_DIR: $!";    
    $DYLD_BEING_REPLACED = 1;
}

foreach my $tool (@TOOLS_NEEDED) {
    if ($DYLD_BEING_REPLACED || $transcript =~ m/^(. .|h|l) \Q$prefix$tool/m) {
        (system("ditto", $tool, "$TOOL_DIR/" . basename($tool)) == 0) or die "Can't ditto $tool to $TOOL_DIR: $!";
    }
}

# Embedded frameworks won't be found by dyld unless they are linked to TOOL_DIR
if (opendir(TDIR, $TOOL_DIR)) {
    while (my $toplevel = readdir(TDIR)) {
        if ($toplevel =~ /\.framework$/ && opendir(FDIR, "$TOOL_DIR/$toplevel/Frameworks")) {
            map { symlink "$toplevel/Frameworks/$_", "$TOOL_DIR/$_" } readdir(FDIR);
            close FDIR;
        }
    }
    close TDIR;
}

# Make sure libraries can be found by any install name
foreach my $binary (keys %links_to_create) {
    -e "$TOOL_DIR/$links_to_create{$binary}" and symlink "$links_to_create{$binary}", "$TOOL_DIR/$binary";
}

print STDERR strftime('%+', localtime) . " Finished copying tools";
$ENV{'DYLD_PRINT_LIBRARIES'} = 1;

print "%ENDPOLE" if ($use_ihook);

# Finally try to apply the transcript
print "Applying system updates...";
$return_value = system("lapply $LAPPLY_FLAGS") >> 8;

if ($return_value == 0) {
    print "%BEGINPOLE" if ($use_ihook);
    print "Running post-actions...";

    # Clear caches (or request rebuild) if it looks like we invalidated them
    unlink glob "./Library/Caches/com.apple.LaunchServices*"    if ($clear_ls_cache);
    unlink glob "./System/Library/Caches/com.apple.Components*" if ($clear_components_cache);
    system_safe "touch ./System/Library/Extensions"             if ($rebuild_kernel_caches);

    # Datestamp the loginwindow
    my $datestamp = "updated " . strftime('%e %B %Y', localtime);
    $datestamp =~ s/  / /;
    system_safe("PlistBuddy", "-c", "set :LoginwindowText '$datestamp'", "./Library/Preferences/com.apple.loginwindow.plist");
} else {
    system_safe("PlistBuddy", "-c", "set :LoginwindowText '$last_datestamp'", "./Library/Preferences/com.apple.loginwindow.plist") if ($last_datestamp);
}

if ($rebuild_kernel_caches || -e $TOOL_DIR) {
    unless ($reboot_after++) {
        for (0..9) {
            printf "Important system files were replaced. This machine will reboot %d seconds.\n", (10-$_);
            sleep 1;
        }
    }
}

if ($reboot_after) {
    print "Rebooting now.";
    system_safe "reboot";
}

cleanup;
exit;