Fix TTY implementation

[libremanage.git] / libremanage.py

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

Copyright (C) 2018 Alyssa Rosenzweig

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

Server names are defined in the accompanying config.py.

Valid commands are as follows:

    - shutdown, reboot, poweron: Power management
    - tty: Open TTY in GNU Screen
"""

import sys
import json
import functools
import subprocess

def open_ssh(server, command):
    config = server["ssh"]
    subprocess.run(["ssh", "-t", config["username"] + "@" + config["host"], "-p", str(config["port"]), command])

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

def set_server_power(state, server):
    print("Setting to power state " + str(state))
    conf = server["power"]

    # Set invert to write LOW for power on and HIGH for off
    if conf["invert"]:
        state = 1 - state

    print(conf["pin"])

    # Export pin, configure, write value, unexport
    open_ssh(server, "echo " + str(conf["pin"]) + " > /sys/class/gpio/export")
    open_ssh(server, "echo out > /sys/class/gpio/gpio" + str(conf["pin"]) + "/direction")
    open_ssh(server, "echo " + str(state) + " > /sys/class/gpio/gpio" + str(conf["pin"]) + "/value")
    open_ssh(server, "echo " + str(conf["pin"]) + " > /sys/class/gpio/unexport")

COMMANDS = {
        # Power managemment

        "shutdown": functools.partial(set_server_power, 0),
        "poweron": functools.partial(set_server_power, 1),
        "reboot": lambda s: (set_server_power(0, s), set_server_power(1, s)),

        # TTY access (or keyboard if wired as such)

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

        # SSH sanity tests

        "sanity": lambda s: open_ssh(s, "whoami"),
        "console": lambda s: open_ssh(s, ""),
}

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

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

    callback(server)

# Load configuration, get command, and go!

with open("config.json") as f:
    CONFIG = json.load(f)

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

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