diff --git a/vm.py b/vm.py index db2bcd179bb09245db48c6801589126d2877adad..c99d3afdc892456543f54ad930579e479c9b4a3a 100755 --- a/vm.py +++ b/vm.py @@ -48,19 +48,21 @@ S=S() force=False no_daemon=False +verbose=1 # daemon is verbose -def r(*arg): - print(">", " ".join(arg)) - run(arg, check=not force) +def r(*arg, check=None, stdin=None): + if check is None: + check = not force + if verbose: print(">", " ".join(arg)) + if stdin is None: + run(arg, check=check) + else: + run(arg, check=check, input=stdin) def nft(rules): - print("\n".join("@ "+i for i in rules.split("\n"))) - #p = subprocess.Popen(["nft", "-i"], stdin=PIPE, encoding='utf-8') + if verbose: print("\n".join("@ "+i for i in rules.split("\n"))) run(["nft", rules], check=not force) - #p.communicate(input=rules) - #p.wait() - #if p.returncode: raise RuntimeError("Wrong returnoce") parser = argparse.ArgumentParser() @@ -71,6 +73,7 @@ subcommands = {} parser.add_argument("-f", "--force", action='store_true') parser.add_argument("-r", "--root_folder", type=str) +parser.add_argument("-v", "--verbose", action='count') def get_spec(f): import inspect @@ -78,7 +81,8 @@ def get_spec(f): f.spec = inspect.getfullargspec(f) return f.spec -def cmd(f): +def cmd(f, ): + if f is None: return f import inspect spec = get_spec(f) subcommands[f.__name__] = f @@ -86,16 +90,26 @@ def cmd(f): for i, arg in enumerate(spec.args): annotation = spec.annotations.get(arg, None) has_default = spec.defaults is not None and i >= len(spec.args) - len(spec.defaults) + if has_default: + default = spec.defaults[i - len(spec.args) + len(spec.defaults)] if annotation in [str, int, float]: f.parser.add_argument( ("--" if has_default else "")+arg, type=annotation, ) if annotation in [bool]: - f.parser.add_argument( - "--"+arg, - action="store_true", - ) + if has_default and default is True: + f.parser.add_argument( + "--no_"+arg, + action="store_false", + dest=arg, + default=True, + ) + else: + f.parser.add_argument( + "--"+arg, + action="store_true", + ) if spec.varargs is not None: arg = spec.varargs annotation = spec.annotations.get(arg, None) @@ -136,30 +150,47 @@ def ask_server(in_struct): import socket connection = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) connection.connect(root_folder+socket_path) + if verbose: print("ASK", in_struct) in_data = json.dumps(in_struct).encode('utf-8') connection.sendall(in_data) connection.shutdown(socket.SHUT_WR) out_data = recvall(connection) out_struct = json.loads(out_data) + if verbose: print("->", out_struct) return out_struct -def daemon(f): - spec = get_spec(f) - assert spec.args[0] == 'ucred' - spec = spec._replace(args=spec.args[1:]) - daemon_funcs[f.__name__] = f - if is_daemon: - return f - # TODO validate types - def l(*arg, **kvarg): - if no_daemon: - f(my_ucred(), *arg, **kvarg) - r = ask_server({"fname":f.__name__, "arg": arg, "kvarg": kvarg}) - return r["return"] +def daemon(root_only=False): + def ll(f): + spec = get_spec(f) + assert spec.args[0] == 'ucred' + spec = spec._replace(args=spec.args[1:]) + + if root_only: + def l(ucred, *arg, **kvarg): + assert ucred.uid == 0 + f(ucred, *arg, **kvarg) + l.__name__ = f.__name__ + l.spec = get_spec(f) + daemon_funcs[f.__name__] = l + else: + daemon_funcs[f.__name__] = f + + if is_daemon: + return f - l.__name__ = f.__name__ - l.spec = spec - return l + # TODO validate types + if root_only and my_ucred().uid != 0: + return None + def l(*arg, **kvarg): + if no_daemon: + f(my_ucred(), *arg, **kvarg) + r = ask_server({"fname":f.__name__, "arg": arg, "kvarg": kvarg}) + return r["return"] + + l.__name__ = f.__name__ + l.spec = spec + return l + return ll ########################################################## @@ -172,6 +203,7 @@ state_not_registered = "not registered" state_powered_off = "powered off" state_aborted = "aborted" state_saved = "saved" +state_running = "running" def vm_dir(vm: str): return f"{root_folder}/{vm}.vm/" @@ -190,7 +222,7 @@ def name_to_id(name: str) -> str: @cmd def name(vm: str) -> str: vm = name_to_id(vm) - return open(vm+".vm/name").read().strip() + return open(vm_dir(vm)+"name").read().strip() def is_valid_name(name): return all(i.isalpha() or i.isnumeric() or i in "-_" for i in name) and any(i.isalpha() for i in name) @@ -198,10 +230,12 @@ def is_valid_name(name): def is_valid_id(id): return all(i.isnumeric() for i in id) +@daemon() def has_read_acces(ucred, vm: str): # TODO! return True +@daemon() def has_write_acces(ucred, vm: str): # TODO! return True @@ -214,10 +248,13 @@ def get_ip(vm: str) -> str: @cmd def get_permanency(vm: str): - return open(vm_dir(vm)+"permanency").read().strip() + try: + return open(vm_dir(vm)+"permanency").read().strip() + except FileNotFoundError: + return "undef" @cmd -@daemon +@daemon() def set_permanency(ucred, vm: str, permanency: str): vm = name_to_id(vm) assert has_write_acces(ucred, vm) @@ -234,32 +271,29 @@ def give_to_user(vm: str, uid: int, gid: Optional[int] = None): os.chown(vm_dir(vm)+"id_ed25519", uid, gid) @cmd -@daemon +@daemon() def state(ucred, vm: str): vm = name_to_id(vm) assert has_read_acces(ucred, vm) - vm = name_to_id(vm) p = run(["VBoxManage", "showvminfo", vm], capture_output=True, encoding='utf8') if "Could not find a registered machine named " in p.stderr: return state_not_registered for l in p.stdout.split('\n'): if l.startswith("State:"): return l[10:].split("(")[0].strip() - print(p.stderr, file=sys.stderr) - raise RuntimeError() + raise RuntimeError(p.stderr) @cmd -def create_from_img(target: str, new_ssh: bool =True, target_name = None): +@daemon(root_only=True) +def create_from_img(ucred: Ucred, target: str, new_ssh: bool = True, target_name = None): target_name = target_name or target target = name_to_id(target) target_dir = vm_dir(target) vnc_passwd = random_passwd() - with open(target_dir+"name", "w") as f: f.write(target_name) - with open(target_dir+"permanency", "w") as f: f.write(f"init {int(time.time())}") with open(target_dir+"vnc_passwd", "w") as f: f.write(vnc_passwd) @@ -318,28 +352,26 @@ def create_from_img(target: str, new_ssh: bool =True, target_name = None): r('VBoxManage', 'internalcommands', 'sethduuid', target_dir+"disk.vmdk") - r('VBoxManage', 'createvm', f'--name={target}', f"--basefolder={folder}/{target}.vm", "--register") + r('VBoxManage', 'createvm', f'--name={target}', f"--basefolder={vm_dir(target)}", "--register") r('VBoxManage', 'modifyvm', target, '--firmware=efi') r('VBoxManage', 'modifyvm', target, '--memory=4096') r("VBoxManage", "storagectl", target, '--name', "SATA Controller", '--add', 'sata', '--bootable', 'on') - r("VBoxManage", "storageattach", target, "--storagectl", "SATA Controller", "--port", "0", "--device", "0", "--type", "hdd", "--medium", f"{folder}/{target_dir}/disk.vmdk") + r("VBoxManage", "storageattach", target, "--storagectl", "SATA Controller", "--port", "0", "--device", "0", "--type", "hdd", "--medium", f"{vm_dir(target)}/disk.vmdk") - create_net(target) + create_net(my_ucred(), target) @cmd -@daemon +@daemon() def clone(ucred, target: str, base: str = "base") -> str: base = name_to_id(base) assert has_read_acces(ucred, base) - target = clone_copy(target, vm_dir(base)+"img") + target = clone_copy(ucred, target, vm_dir(base)+"img") give_to_user(target, ucred.uid, ucred.gid) return target -@cmd -def clone_copy(target: str, img_path: str) -> str: +def create_vm_dir(target: str) -> str: import random - assert is_valid_name(target) target_id = f"{random.randint(0, 999999):06}" assert not os.path.exists(target_id+".vm") assert not os.path.exists(target+".vm") @@ -348,61 +380,112 @@ def clone_copy(target: str, img_path: str) -> str: os.mkdir(target_dir) r('ln', '-sr', target_id+".vm", target+".vm") - r('cp', '--reflink', "-n", img_path, target_dir+"img") - create_from_img(target_id, target_name=target) + with open(vm_dir(target)+"name", "w") as f: f.write(target) + return target_id + +@cmd +@daemon(root_only=True) +def clone_copy(ucred, target: str, img_path: str) -> str: + assert is_valid_name(target) + + target_id = create_vm_dir(target) + r('cp', '--reflink', "-n", img_path, vm_dir(target_id)+"img") + create_from_img(ucred, target_id, target_name=target) return target_id def ssh_args(vm: str, *arg: tuple[str, ...], user: str = "u"): vm, user = extended_name(vm, user) vm = name_to_id(vm) - return ['ssh', "-i", vm_dir(vm)+"id_ed25519", + arg = ['ssh', "-i", vm_dir(vm)+"id_ed25519", "-o", f"UserKnownHostsFile={vm_dir(vm)}/known_hosts", "-o", "HostKeyAlgorithms=ssh-ed25519", "-o", f"HostKeyAlias=vm_{vm}", f"{user}@{get_ip(vm)}", *arg] + if verbose: print(">>", " ".join(arg)) + return arg + @cmd def ssh(vm: str, *arg: tuple[str, ...], user: str = "u"): run(ssh_args(vm, *arg, user=user)) +sshfs_root = lambda: os.environ["HOME"]+f"/m/vm/" + +@cmd +def sshfs_mountdir(vm: str, user: str = "u"): + return sshfs_root()+f"/{user}@{name(vm)}" + +@cmd +def sshfs(vm: str, user: str = None): + vm, user = extended_name(vm, user) + if user is None: + sshfs(vm, "root") + sshfs(vm, "u") + return + mount_dir = sshfs_mountdir(vm, user) + if os.path.isdir(mount_dir) and len(os.listdir(mount_dir)) != 0: + return + r("mkdir", "-p", mount_dir) + r("sshfs", f"{user}@{get_ip(vm)}:/", mount_dir, "-o", f"ssh_command=ssh -i {vm_dir(vm)}/id_ed25519 -o UserKnownHostsFile={vm_dir(vm)}/known_hosts -o HostKeyAlgorithms=ssh-ed25519 -o HostKeyAlias=vm_{vm}") + if not os.path.islink(mount_dir+'~'): + home_dir = "/root" if user == "root" else f"/home/{user}" + r("ln", "-sr", mount_dir+home_dir, mount_dir+"~") + +@cmd +def sshfs_clean(): + root = sshfs_root() + for f in os.listdir(root): + if os.path.isdir(root+f) and len(os.listdir(root+f)) == 0: + r("rmdir", root+f) + for f in os.listdir(root): + if os.path.islink(root+f) and not pathlib.Path(root+f+"/").absolute().exists(): + r("rm", root+f) + + +def escape_sh(s): + return "'" + s.replace("'", "'\"'\"'") + "'" @cmd def vncapp(vm: str, cmd: str): import random import psutil + unit_id = random.randint(100000, 999999) vm, user = extended_name(vm) vm = name_to_id(vm) - display_id=random.randint(1, 50) - vnc_server = subprocess.Popen(ssh_args(vm, f"(cat /vnc_passwd;echo; cat /vnc_passwd; echo;echo n) | vncpasswd; vncserver :{display_id}", user=user)) + display_id=random.randint(10, 50) + vnc_server = subprocess.Popen(ssh_args(vm, f"systemd-run --unit vncapp-vnc-{display_id}-{unit_id} --user -P bash -c '(cat /vnc_passwd;echo; cat /vnc_passwd; echo;echo n) | vncpasswd; vncserver :{display_id}'", user=user)) vnc_client_env = os.environ.copy() vnc_client_env["VNC_PASSWORD"] = open(vm_dir(vm)+"vnc_passwd", "r").read().strip() - app = subprocess.Popen(ssh_args(vm, f"DISPLAY=:{display_id}", cmd, user=user)); + app = subprocess.Popen(ssh_args(vm, f"systemd-run --unit vncapp-app-{display_id}-{unit_id} --user -P -E DISPLAY=:{display_id} bash -c {escape_sh(cmd)}", user=user)); time.sleep(1) vnc_client = subprocess.Popen(["vncviewer", get_ip(vm)+f":{display_id}"], env=vnc_client_env) def on_terminate(proc): - print("KILLING ALL APPS") + if verbose: print("KILLING ALL APPS") vnc_server.send_signal(15) vnc_client.send_signal(15) app.send_signal(15) psutil.wait_procs([vnc_client, app, vnc_server], callback=on_terminate) + ssh(vm, f"systemctl --user stop vncapp-vnc-{display_id}-{unit_id} vncapp-app-{display_id}-{unit_id}") @cmd def vncsession(vm: str, display_id: int =0): import random import psutil + unit_id = random.randint(100000, 999999) vm, user = extended_name(vm) vm = name_to_id(vm) - vnc_server = subprocess.Popen(ssh_args(vm, f"(cat /vnc_passwd;echo; cat /vnc_passwd; echo;echo n) | vncpasswd; vncserver :{display_id}", user=user)) + vnc_server = subprocess.Popen(ssh_args(vm, f"systemd-run --unit vncsession-{display_id}-{unit_id} --user -P bash -c '(cat /vnc_passwd;echo; cat /vnc_passwd; echo;echo n) | vncpasswd; vncserver :{display_id}'", user=user)) vnc_client_env = os.environ.copy() vnc_client_env["VNC_PASSWORD"] = open(vm_dir(vm)+"vnc_passwd", "r").read().strip() time.sleep(1) vnc_client = subprocess.Popen(["vncviewer", get_ip(vm)+f":{display_id}"], env=vnc_client_env) def on_terminate(proc): - print("KILLING ALL APPS") + if verbose: print("KILLING ALL APPS") vnc_server.send_signal(15) vnc_client.send_signal(15) psutil.wait_procs([vnc_client, vnc_server], callback=on_terminate) + ssh(vm, f"systemctl --user stop vncsession-{display_id}-{unit_id}") @@ -413,9 +496,38 @@ def all_virtuals(): if is_valid_id(vm): yield vm +def terminal_len(val: str) -> int: + return len(val) + +def format_table(table): + lengths = [] + for r in table: + for i, c in enumerate(r): + if len(lengths) <= i: + lengths.append(0) + lengths[i] = max(lengths[i], terminal_len(c)) + + for r in table: + print(" ".join(c + " "*(clen-terminal_len(c)) for c, clen in zip(r, lengths))) + +@cmd +def index(color: bool = True): + out = [] + for vm in all_virtuals(): + out_state = state(vm) + if out_state in [state_running]: + if try_ping(vm): + out_state += " (pinging)" + else: + out_state += " (NO PING)" + out_rw = ('w' if has_write_acces(vm) else 'r') if has_read_acces(vm) else '-' + out.append([vm, out_rw, name(vm), out_state, get_permanency(vm)]) + return format_table(out) + + @cmd -@daemon +@daemon() def clean(ucred): for vm in all_virtuals(): if has_write_acces(ucred, vm): @@ -430,7 +542,10 @@ def clean(ucred): @cmd -def remove_net(vm: str): +@daemon() +def remove_net(ucred, vm: str): + vm = name_to_id(vm) + assert has_write_acces(ucred, vm) network_dir = vm_dir(vm)+"network/" if os.path.isfile(network_dir+"interface"): interface = open(network_dir+"interface").read().strip() @@ -438,7 +553,7 @@ def remove_net(vm: str): interface = open(network_dir+"interface").read().strip() if open(network_dir+"boot_id", "r").read().strip() == boot_id: if os.path.isfile(network_dir+"dhcp.pid"): - r("kill", open(network_dir+"dhcp.pid").read().strip()) + r("kill", open(network_dir+"dhcp.pid").read().strip(), check=False) net_id = int(open(network_dir+"net_id", "r").read()) r("ip", "addr", "del", f"10.37.{net_id}.1/24", "dev", interface) @@ -447,8 +562,12 @@ def remove_net(vm: str): r('rm', '-r', network_dir) @cmd -def create_net(vm: str): - remove_net(vm) +@daemon() +def create_net(ucred, vm: str): + vm = name_to_id(vm) + assert has_write_acces(ucred, vm) + + remove_net(ucred, vm) network_dir = vm_dir(vm)+"network/" os.mkdir(network_dir) @@ -466,7 +585,7 @@ def create_net(vm: str): net_id = int(interface[7:]) - print("interface", interface) + if verbose: print("interface", interface) r("VBoxManage", "hostonlyif", "ipconfig", interface, f"--ip=10.37.{net_id}.1", "--netmask=255.255.255.0") @@ -510,19 +629,18 @@ def create_net(vm: str): r("dhcpd", "-4", "-cf", network_dir+"dhcp.config", "-pf", network_dir+"dhcp.pid", "-lf", network_dir+"dhcp.lp", interface) @cmd -@daemon -def modify_net(ucred, vm: str, van: bool = False, lan: bool = False, pc: bool = False, pc_all: bool = False): +@daemon() +def modify_net(ucred, vm: str, wan: bool = False, lan: bool = False, pc: bool = False, pc_all: bool = False): vm = name_to_id(vm) assert has_write_acces(ucred, vm) assert not (pc_all and not pc) - assert not (lan and not van) network_dir = vm_dir(vm)+"network/" interface = open(network_dir+"interface").read().strip() net_id = open(network_dir+"net_id").read().strip() assert open(network_dir+"boot_id", "r").read().strip() == boot_id local_ips = "{10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16}" - if van: + if wan: pass todo = [f"flush chain inet filter input_from_{interface}", f"flush chain inet filter forward_from_{interface}", @@ -539,9 +657,9 @@ def modify_net(ucred, vm: str, van: bool = False, lan: bool = False, pc: bool = todo.append(f"add rule inet filter forward_to_{interface} ip saddr {local_ips} drop") else: todo.append(f"add rule inet filter forward_from_{interface} ip daddr {local_ips} drop") - todo.append(f"add rule inet filter forward_from_{interface} ip saddr {local_ips} drop") + todo.append(f"add rule inet filter forward_to_{interface} ip saddr {local_ips} drop") - if van: + if wan: todo.append(f"add rule inet filter forward_from_{interface} accept") todo.append(f"add rule inet filter forward_to_{interface} ct state {{ established, related }} accept") @@ -552,16 +670,18 @@ def modify_net(ucred, vm: str, van: bool = False, lan: bool = False, pc: bool = nft("\n".join(todo)) @cmd -@daemon +@daemon() def start(ucred, vm: str): vm = name_to_id(vm) assert has_write_acces(ucred, vm) + if open(vm_dir(vm)+"network/boot_id", "r").read().strip() != boot_id: + create_net(ucred, vm) r("VBoxManage", "startvm", vm, "--type=headless") if get_permanency(vm).startswith("init "): set_permanency(ucred, vm, "tmp") @cmd -@daemon +@daemon() def try_ping(ucred, vm: str) -> bool: vm = name_to_id(vm) assert has_write_acces(ucred, vm) @@ -582,7 +702,8 @@ def start_and_wait(vm: str): start(vm) wait_started(vm) -@daemon +@cmd +@daemon() def poweroff(ucred, vm: str): vm = name_to_id(vm) assert has_write_acces(ucred, vm) @@ -601,8 +722,20 @@ def poweroff_and_wait(vm: str): time.sleep(1) return +def poweroff_and_wait_daemon(ucred, vm: str): + vm = name_to_id(vm) + if state(ucred, vm) in [state_powered_off, state_not_registered, state_aborted, state_saved]: return + poweroff(ucred, vm) + while True: + time.sleep(0.1) + s = state(ucred, vm) + print("state", s) + if s == state_powered_off: + time.sleep(1) + return + @cmd -@daemon +@daemon() def get_tmp_name(ucred) -> str: last_id = int(open("last_tmp_id").read().strip()) last_id = (last_id+1) % 1000 @@ -615,7 +748,7 @@ if is_daemon: prepared_forks = {} @cmd -@daemon +@daemon() def get_prepared_fork(ucred, base: str = "base") -> Optional[str]: name_to_id(base) assert has_read_acces(ucred, base) @@ -626,9 +759,8 @@ def get_prepared_fork(ucred, base: str = "base") -> Optional[str]: return vm @cmd -@daemon +@daemon(root_only=True) def prepare_forks(ucred, base: str = "base", count: int =1) -> Optional[str]: - assert ucred.uid == 0 if not base in prepared_forks: prepared_forks[base] = [] if len(prepared_forks[base]) < count: @@ -653,18 +785,30 @@ def extended_name(vm: str, user: str = "u") -> tuple[str, str]: assert not is_daemon if len(vm.split("@"))==2: user, vm = vm.split("@") - if len(vm.split("~"))==2: - base, vm = vm.split("~") + + do_power_on = False + if len(vm.split("!"))==2: + vm, tmp = vm.split("!") + do_power_on = True + assert tmp == "" + + if len(vm.split("+"))==2: + base, vm = vm.split("+") if not vm : vm = get_tmp_vm(base or "base") else: vm = clone(vm, base or "base") start(vm) wait_started(vm) + vm = name_to_id(vm) + + if do_power_on: + if state(vm) != state_running: + start_and_wait(vm) return vm, user -@daemon +@daemon() def remove_force(ucred, vm: str, keep_image: bool = False): vm = name_to_id(vm) if os.path.isfile(f"{vm}.vm/no_remove"): @@ -672,7 +816,7 @@ def remove_force(ucred, vm: str, keep_image: bool = False): assert has_write_acces(ucred, vm) if state(ucred, vm) != state_not_registered: r("VBoxManage", "unregistervm", vm, "--delete-all") - remove_net(vm) + remove_net(ucred, vm) if keep_image: for f in os.listdir(vm+".vm"): if f != 'img': @@ -689,10 +833,18 @@ def remove(vm: str, keep_image: bool = False): remove_force(vm, keep_image=keep_image) +@cmd +@daemon(root_only=True) +def exit_server(ucred): + for vm in all_virtuals(): + poweroff_and_wait_daemon(ucred, vm) + remove_net(ucred, vm) + exit(0) + ########################################################## -if is_daemon: +def main_daemon(): import socket import struct try: @@ -704,6 +856,7 @@ if is_daemon: server.bind(socket_path) os.chmod(socket_path, 0o777) + import traceback try: while True: @@ -723,7 +876,6 @@ if is_daemon: try: res = f(ucred, *in_struct["arg"], **in_struct["kvarg"]) except Exception as e: - import traceback traceback.print_exception(e) out_struct = {'excepttion': str(type(e))} else: @@ -731,24 +883,36 @@ if is_daemon: print("OUT", out_struct) out_data = json.dumps(out_struct).encode('utf-8') print("OUT", out_data) - connection.sendall(out_data) - connection.close() + sys.stdout.flush() + sys.stderr.flush() + try: + connection.sendall(out_data) + except Exception as e: + traceback.print_exception(e) + try: + connection.close() + except Exception as e: + traceback.print_exception(e) finally: # close the connection # remove the socket file os.unlink(socket_path) exit(1) -if __name__ == "__main__": +def main(): args = parser.parse_args() - print(args) + + global verbose + verbose = args.verbose + force = args.force + if args.root_folder is not None: + root_folder = args.root_folder+"/" + + if verbose: print(args) if not args.subcommand: parser.print_help() else: import inspect - force = args.force - if args.root_folder is not None: - root_folder = args.root_folder+"/" f = subcommands[args.subcommand] spec = get_spec(f) f_kvarg = {} @@ -766,9 +930,19 @@ if __name__ == "__main__": if annotation == tuple[str, ...]: f_arg += args.__dict__[arg] - print(f_arg, f_kvarg) + if verbose: print(f_arg, f_kvarg) r = f(*f_arg, **f_kvarg) - if r is not None: print(r) + if r is not None: + if isinstance(r, tuple) or isinstance(r, list): + for i in r: + print(i) + else: + print(r) +if __name__ == "__main__": + if is_daemon: + main_daemon() + else: + main()