libremanage/libremanage

#!/usr/bin/python3

"""
libremanage - Lightweight, free software for remote side-chanel server management

Copyright (C) 2018 Alyssa Rosenzweig <alyssa@rosenzweig.io>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as published by
the Free Software Foundation, either version 3 of the License, 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 Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

USAGE = """
Usage:

    $ libremanage [server name] [command]

Example:

    $ libremanage web2 reboot

Valid commands are as follows:

    - shutdown, reboot, poweron: Power management
    - tty: Open TTY in GNU Screen
    - sanity: Sanity test to check if manager is reachable
    - sanity-sh: Sanity test exposing a shell on the manager

Define a configuration file in ~/.libremanage.json. See the included
config.json for an example. Named servers correspond to managed servers;
managers correspond to single-board computers connecting the servers.
libremanage SSHs into the manager to access the server through the
side-channel.
"""

import sys
import json
import functools
import subprocess
import time
import os.path

def open_ssh(server, command, force_tty=False):
    cfg = server["ssh"]
    args = ["ssh"] + (["-t"] if force_tty else []) + [cfg["username"] + "@" + cfg["host"], "-p", str(cfg["port"]), command]
    subprocess.run(args)

def die_with_usage(message):
    print(message)
    print(USAGE)
    sys.exit(1)

def get_server_handle(name):
    try:
        server = CONFIG["servers"][name]
    except KeyError:
        die_with_usage("Unknown server, please configure")

    # Associate manager configuration
    server["ssh"] = CONFIG["managers"][server["manager"]]

    return server

"""
Power management: currently, we only support the `hidusb-relay-cmd` driver,
wired up as to the power button pins. We may want to expose more options in the
config for other boards.
"""

POWER_OFF    = 0
POWER_ON     = 1
POWER_REBOOT = 2

def set_server_power(state, server):
    conf = server["power"]

    # Ensure the button is in a known state
    power_write(server, conf, 0)

    if state == POWER_OFF or state == POWER_ON:
        power_button(server, conf, state)
    elif state == POWER_REBOOT:
        # Requires that we already be online.
        power_button(server, conf, POWER_OFF)
        power_button(server, conf, POWER_ON)

def power_write(server, conf, state):
    if conf["type"] == "hidusb-relay-cmd":
        verb = "on" if state == 1 else "off"
        open_ssh(server, "hidusb-relay-cmd ID=" + conf["relay"] + " " + verb + " " + str(conf["channel"]))
    else:
        die_with_usage("Unknown power type " + conf["type"])

def power_button(server, conf, state):
    # Hold down the power to force off (via the EC),
    # or just flick on to turn on

    power_write(server, conf, 1)
    time.sleep(conf["timing"]["off" if state == POWER_OFF else "on"])
    power_write(server, conf, 0)

"""
Define the list of commands implemented as a dict mapping names to functions
actuating the command
"""

COMMANDS = {
        # Power managemment

        "shutdown": functools.partial(set_server_power, POWER_OFF),
        "poweron": functools.partial(set_server_power, POWER_ON),
        "reboot": functools.partial(set_server_power, POWER_REBOOT),

        # TTY access (or keyboard if wired as such)

        "tty": lambda s: open_ssh(s, "screen " + s["tty"]["file"] + " " + str(s["tty"]["baud"]), force_tty=True),

        # SSH sanity tests

        "sanity": lambda s: open_ssh(s, "whoami"),
        "sanity-sh": lambda s: open_ssh(s, ""),
}

def issue_command(server_name, command):
    server = get_server_handle(server_name)

    try:
        callback = COMMANDS[command]
    except KeyError:
        die_with_usage("Invalid command supplied")

    callback(server)

# Load configuration, get command, and go!

try:
    with open(os.path.expanduser("~/.libremanage.json")) as f:
        CONFIG = json.load(f)
except FileNotFoundError:
    die_with_usage("Configuration file missing in ~/.libremanage.json")

if len(sys.argv) != 3:
    die_with_usage("Incorrect number of arguments")

issue_command(sys.argv[1], sys.argv[2])