pbot/applets/pbot-vm/host/bin/vm-exec

#!/usr/bin/env perl

# File: vm-exec
#
# Purpose: Process and send commands to the PBot Guest server (guest-server)
# using the details from the config/vm-exec.json configuration file.
#
# Additionally, takes `-revert` and `-health` options to revert VM or check
# VM's health.
#
# Use the PBOTVM_CID, PBOTVM_VPORT and/or PBOTVM_SERIAL environment variables to
# override the config/vm-exec.json values. E.g.:
#
# $ PBOTVM_CID=42 PBOTVM_SERIAL=7777 vm-exec -lang=sh echo test

# SPDX-FileCopyrightText: 2021-2024 Pragmatic Software <pragma78@gmail.com>
# SPDX-License-Identifier: MIT

use 5.020;

use warnings;
use strict;

use feature qw(signatures);
no warnings qw(experimental::signatures);

use constant {
    DOMAIN  => $ENV{PBOTVM_DOMAIN} // 'pbot-vm',
    ADDR    => $ENV{PBOTVM_ADDR}   // '127.0.0.1',
    SERIAL  => $ENV{PBOTVM_SERIAL} // 5555,
    HEALTH  => $ENV{PBOTVM_HEALTH} // 5556,
    CID     => $ENV{PBOTVM_CID}    // 7,
    VPORT   => $ENV{PBOTVM_VPORT}  // 5555,
    VAGRANT => $ENV{PBOTVM_VAGRANT} // 0,
};

use File::Basename;
use JSON::XS;
use IPC::Open2;
use IO::Socket;

use FindBin qw($RealBin);
use lib "$RealBin/../lib";

sub connect_vsock($context) {
    return undef if not $context->{'vm-cid'};

    print STDERR "Connecting to remote VM socket CID $context->{'vm-cid'} port $context->{'vm-vport'}\n";

    my $command = "socat - VSOCK-CONNECT:$context->{'vm-cid'}:$context->{'vm-vport'}";

    my ($pid, $input, $output) = eval {
        my $pid = open2(my $output, my $input, $command);
        return ($pid, $input, $output);
    };

    if ($@) {
        print STDERR "Failed to connect to VM socket: $@\n";
        return undef;
    }

    if (not defined $pid) {
        print STDERR "Failed to connect to VM socket: $!\n";
        return undef;
    }

    return ($input, $output);
}

sub connect_serial($context) {
    print STDERR "Connecting to remote VM serial port $context->{'vm-serial'}\n";

    my $vm = IO::Socket::INET->new(
        PeerAddr => $context->{'vm-addr'},
        PeerPort => $context->{'vm-serial'},
        Proto => 'tcp',
        Type => SOCK_STREAM
    );

    # return same $vm handle for ($input, $output)
    return ($vm, $vm);
}

sub connect_vm($context) {
    my ($input, $output);

    # attempt preferred VSOCK connection
    ($input, $output) = connect_vsock($context);

    # fallback to serial
    if (not defined $input) {
        ($input, $output) = connect_serial($context);
    }

    if (not defined $input) {
        die "Could not create connection to VM: $!";
    }

    print STDERR "Connected to VM.\n";
    return ($input, $output);
}

sub make_context_from_args(@args_in) {
    my $args = join ' ', @args_in;

    # extract leading options
    my %opts;
    while ($args =~ s/^-(revert|health)\s+//) {
        $opts{$1} = 1;
    }

    my $context = eval { decode_json $args };

    if ($@) {
        # wasn't JSON; make structure manually
        $context = { code => $args };

        # command-line usage
        if (not length $context->{code}) {
            die "Usage: $0 [-lang=<language>] <code>\n";
        }
    }

    # set extracted leading options
    foreach my $opt (keys %opts) {
        print STDERR "Setting option `$opt`.\n";
        $context->{$opt} = 1;
    }

    # parse options specific to vm-exec
    while ($context->{code} =~ s/^-(lang|revert|health|vm-domain|vm-health|vm-cid|vm-vport|vm-serial|vm)=([^ ]+)\s*//) {
        my ($option, $value) = ($1, $2);
        print STDERR "Overriding `$option` to `$value`.\n";
        $context->{$option} = lc $value;
    }

    return $context;
}

sub load_language($context) {
    my $language = $context->{lang};

    eval {
        require "Languages/$language.pm";
    } or do {
        my @languages;

        foreach my $module (sort glob "$RealBin/../lib/Languages/*.pm") {
            $module = basename $module;
            next if $module =~ m/^_/;
            $module =~ s/.pm$//;

            require "Languages/$module.pm" or die $!;

            my $mod = "Languages::$module"->new(%$context);

            if (exists $mod->{name} && $mod->{name} eq $language) {
                return $mod;
            }

            $module = $mod->{name} if exists $mod->{name};
            push @languages, $module;
        }

        print "Language '$language' is not supported.\nSupported languages are: ", join(', ', @languages), "\n";
        exit 1;
    };

    return "Languages::$language"->new(%$context);
}

sub load_config($file) {
    open my $fh, '<', $file or die "load config: could not open $file: $!\n";
    local $/ = undef;
    my $json = <$fh>;
    close $fh;
    return decode_json($json);
}

sub config_get_machine($config, $machine) {
    # check if machine is an alias and update to actual machine
    if (exists $config->{aliases}->{$machine}) {
        $machine = $config->{aliases}->{$machine};
    }

    # return pointer to config entry
    return $config->{machines}->{$machine};
}

sub list_machines($config) {
    my @machines;

    # quick-and-easy: just list the aliases
    foreach my $alias (keys %{$config->{aliases}}) {
        push @machines, $alias;
    }

    return join ', ', sort @machines;
}

sub configure_context($context, $config) {
    if (exists $context->{vm}) {
        # -vm option was passed, forcibly override machine
        my $machine = $context->{vm};

        # point at config entry
        my $entry = config_get_machine($config, $machine);

        if (not defined $entry) {
            my $machines = list_machines($config);
            print "Unknown machine '$machine'; available machines are: $machines\n";
            exit 3;
        }

        # override values
        $context->{'vm-domain'}  = $machine;
        $context->{'vm-addr'}    = $entry->{'addr'};
        $context->{'vm-health'}  = $entry->{'health'};
        $context->{'vm-serial'}  = $entry->{'serial'};
        $context->{'vm-cid'}     = $entry->{'cid'};
        $context->{'vm-vport'}   = $entry->{'vport'};
        $context->{'vm-vagrant'} = $entry->{'vagrant'};
    } else {
        # otherwise configure any undefined values as default machine
        my $machine = $config->{'default-machine'};

        # point at config entry
        my $entry = config_get_machine($config, $machine);

        if (not defined $entry) {
            my $machines = list_machines($config);
            print "Unknown machine '$machine'; available machines are: $machines\n";
            exit 3;
        }

        # update any undefined values, preserving any existing values
        $context->{'vm-domain'}  //= $machine;
        $context->{'vm-addr'}    //= $entry->{'addr'};
        $context->{'vm-health'}  //= $entry->{'health'};
        $context->{'vm-serial'}  //= $entry->{'serial'};
        $context->{'vm-cid'}     //= $entry->{'cid'};
        $context->{'vm-vport'}   //= $entry->{'vport'};
        $context->{'vm-vagrant'} //= $entry->{'vagrant'};
    }

    # set any undefined values to default values
    $context->{nick}         //= 'vm';
    $context->{channel}      //= 'vm';
    $context->{lang}         //= 'c2x';
    $context->{'vm-domain'}  //= DOMAIN;
    $context->{'vm-addr'}    //= ADDR;
    $context->{'vm-health'}  //= HEALTH;
    $context->{'vm-serial'}  //= SERIAL;
    $context->{'vm-cid'}     //= CID;
    $context->{'vm-vport'}   //= VPORT;
    $context->{'vm-vagrant'} //= VAGRANT;
}

sub main() {
    my $context = make_context_from_args(@ARGV);

    my $config = load_config("$RealBin/../config/vm-exec.json");

    configure_context($context, $config);

    if ($context->{revert} && $context->{health}) {
        print STDERR "-health and -revert cannot be used together; aborting.\n";
        exit 0;
    }

    # instructed to revert machine
    if ($context->{revert}) {
        if (exists $config->{aliases}->{$context->{'vm-domain'}}) {
            $context->{'vm-domain'} = $config->{aliases}->{$context->{'vm-domain'}};
        }

        print STDERR "REVERT $context->{'vm-domain'}\n";

        if ($context->{'vm-vagrant'}) {
            system("virsh -c qemu:///system snapshot-revert $context->{'vm-domain'} 1");
        } else {
            system("virsh snapshot-revert $context->{'vm-domain'} 1");
        }

        exit 0;
    }

    # instructed to check health
    if ($context->{health}) {
        my $health = IO::Socket::INET->new(
            PeerAddr => $context->{'vm-addr'},
            PeerPort => $context->{'vm-health'},
            Proto => 'tcp',
            Type => SOCK_STREAM
        );

        if (not defined $health) {
            print STDERR "Unable to connect to health $context->{'vm-addr'} $context->{'vm-health'}\n";
            exit 2;
        }

        print $health "\n";


        eval {
            alarm 2;
            local $SIG{ALRM} = sub { die "Health timed-out\n"; };

            while (my $output = <$health>) {
                last if $output eq ":END\r\n";
                print $output;
            }
            close $health;
        };

        if ($@) {
            print STDERR "Failed to get health: $@\n";
            exit 1;
        }

        exit 0;
    }

    # load language before checking usage in order to handle -lang=? flag
    # to list languages instead of showing a usage message
    my $lang = load_language($context);

    # now check usage
    if (not length $context->{code}) {
        if (exists $context->{usage}) {
            print "$context->{usage}\n";
        } else {
            print "Usage: cc [-lang=<language>] [-vm=<virtual machine>] [-info] [-paste] [-args \"command-line arguments\"] [compiler/language options] <code> [-stdin <stdin input>]\n";
        }
        exit 1;
    }

    # run the language
    $lang->process_interactive_edit;
    $lang->process_standard_options;
    $lang->process_custom_options;
    $lang->process_cmdline_options;
    $lang->preprocess_code;

    ($lang->{'vm-input'}, $lang->{'vm-output'}) = connect_vm($context);
    $lang->execute;

    $lang->postprocess_output;
    $lang->show_output;
    $lang->done;
}

main();