lxc-network/spawn_network.py

#!/usr/bin/env python3

from lxc_net import parse_network, libvirt_error, util
import signal
import libvirt
import argparse
import logging
import sys


def parse_args():
    parser = argparse.ArgumentParser(description="Spawns a network of LXC containers.")
    parser.add_argument("topology", help="A YAML file defining the network topology")

    args = parser.parse_args()
    return args


def handle_dom(topology, cmd, sudo_drop=True):
    if not cmd:
        print("Missing argument.")
        return

    if cmd[0] in ["up", "down"]:
        dom_id = None
        if len(cmd) > 1:
            try:
                dom_id = [int(x) for x in cmd[1:]]
            except ValueError:
                print("Bad id: {}".format(cmd[1]))
                return
        state = cmd[0] == "up"

        util.ensure_sudo_rights()
        if dom_id:
            for dom in dom_id:
                topology.dom_setstate_single(dom, state, verbose=True)
        else:
            topology.dom_setstate(state, verbose=True)
        if sudo_drop:
            util.drop_sudo_rights()
    elif cmd[0] == "restart":
        handle_dom(topology, ["down"] + cmd[1:], sudo_drop=False)
        handle_dom(topology, ["up"] + cmd[1:])
    elif cmd[0] == "help":
        print("Available commands: up, down, restart, help")
    else:
        print("Bad command.")


def main_loop(topology):
    actions = {
        "dom": handle_dom,
        "exit": "_exit",
        "help": "_help",
    }

    while True:
        print(">> ", end="")
        sys.stdout.flush()

        try:
            line = input().strip().split()
            if not line:  # Empty line
                continue

            action = actions.get(line[0], "_bad")
            if isinstance(action, str):
                if action == "_bad":
                    print("Unknown command: {}.".format(line[0]))
                if action in ["_help", "_bad"]:
                    print("Available commands: dom, exit, help")
                elif action == "_exit":
                    return
                else:
                    print("Unknown action.")
            else:
                action(topology, line[1:])
        except Exception as exn:
            print(exn)


def main():
    args = parse_args()

    received_sigint = False

    def handle_sigint(signum, frame):
        """ Called upon SIGINT (^C) """
        nonlocal received_sigint

        print(" >> Received SIGINT, stopping network...")
        received_sigint = True

    signal.signal(signal.SIGINT, handle_sigint)

    logging.basicConfig(level=logging.WARNING)
    libvirt.registerErrorHandler(libvirt_error.libvirt_error_handler, None)

    conn = libvirt.open("lxc:///")

    topology = parse_network.YamlTopology(args.topology, conn)

    util.ensure_sudo_rights()
    topology.net_start(verbose=True)
    topology.dom_start(verbose=True)
    util.drop_sudo_rights()

    print("Network running. Press ^C to terminate.")

    main_loop(topology)

    util.ensure_sudo_rights()
    topology.dom_stop(verbose=True)
    topology.net_stop(verbose=True)
    util.drop_sudo_rights()


if __name__ == "__main__":
    main()