From a3bcb93f2fa2dd47a180defae7757695fba9790d Mon Sep 17 00:00:00 2001 From: lassulus Date: Mon, 24 Jul 2023 21:11:49 +0200 Subject: [PATCH] clan-cli: split out ssh subcommand, add more tests --- pkgs/clan-cli/clan_cli/__init__.py | 78 ++------------------------- pkgs/clan-cli/clan_cli/ssh.py | 80 ++++++++++++++++++++++++++++ pkgs/clan-cli/tests/test_clan_ssh.py | 57 ++++++++++++++++++++ 3 files changed, 140 insertions(+), 75 deletions(-) create mode 100644 pkgs/clan-cli/clan_cli/ssh.py create mode 100644 pkgs/clan-cli/tests/test_clan_ssh.py diff --git a/pkgs/clan-cli/clan_cli/__init__.py b/pkgs/clan-cli/clan_cli/__init__.py index 4bc36e38..e4ff0dfa 100755 --- a/pkgs/clan-cli/clan_cli/__init__.py +++ b/pkgs/clan-cli/clan_cli/__init__.py @@ -1,10 +1,8 @@ # !/usr/bin/env python3 import argparse -import json -import subprocess import sys -from . import admin +from . import admin, ssh has_argcomplete = True try: @@ -13,64 +11,6 @@ except ImportError: # pragma: no cover has_argcomplete = False -def ssh(args: argparse.Namespace) -> None: - if args.json: - ssh_data = json.load(args.json) - subprocess.run( - [ - "nix", - "shell", - "nixpkgs#sshpass", - "-c", - "torify", - "sshpass", - "-p", - ssh_data.get("password"), - "ssh", - "-o", - "UserKnownHostsFile=/dev/null", - "-o", - "StrictHostKeyChecking=no", - f'root@{ssh_data["address"]}', - ] - + args.ssh_args - ) - elif args.png: - png_text = subprocess.Popen( - [ - "nix", - "shell", - "nixpkgs#zbar", - "-c", - "zbarimg", - "--quiet", - "--raw", - args.png, - ], - stdout=subprocess.PIPE, - ).stdout.read() - ssh_data = json.loads(png_text) - subprocess.run( - [ - "nix", - "shell", - "nixpkgs#sshpass", - "-c", - "torify", - "sshpass", - "-p", - ssh_data.get("password"), - "ssh", - "-o", - "UserKnownHostsFile=/dev/null", - "-o", - "StrictHostKeyChecking=no", - f'root@{ssh_data["address"]}', - ] - + args.ssh_args - ) - - # this will be the entrypoint under /bin/clan (see pyproject.toml config) def main() -> None: parser = argparse.ArgumentParser(description="cLAN tool") @@ -80,19 +20,7 @@ def main() -> None: admin.register_parser(parser_admin) parser_ssh = subparsers.add_parser("ssh", help="ssh to a remote machine") - parser_ssh.add_argument( - "-j", - "--json", - help="specify the json file for ssh data (generated by starting the clan installer", - ) - parser_ssh.add_argument( - "-P", - "--png", - help="specify the json file for ssh data as the qrcode image (generated by starting the clan installer", - ) - # TODO pass all args we don't parse into ssh_args, currently it fails if arg starts with - - parser_ssh.add_argument("ssh_args", nargs="*", default=[]) - parser_ssh.set_defaults(func=ssh) + ssh.register_parser(parser_ssh) if has_argcomplete: argcomplete.autocomplete(parser) @@ -102,7 +30,7 @@ def main() -> None: args = parser.parse_args() if hasattr(args, "func"): - args.func(args) + args.func(args) # pragma: no cover if __name__ == "__main__": # pragma: no cover diff --git a/pkgs/clan-cli/clan_cli/ssh.py b/pkgs/clan-cli/clan_cli/ssh.py new file mode 100644 index 00000000..1771e5df --- /dev/null +++ b/pkgs/clan-cli/clan_cli/ssh.py @@ -0,0 +1,80 @@ +import argparse +import json +import subprocess +from typing import Optional + + +def ssh( + host: str, + user: str = "root", + password: Optional[str] = None, + ssh_args: list[str] = [], +) -> None: + if ssh_args is None: + ssh_args = [] + nix_shell_args = [] + password_args = [] + if password: + nix_shell_args = [ + "nix", + "shell", + "nixpkgs#sshpass", + "-c", + ] + password_args = [ + "sshpass", + "-p", + password, + ] + _ssh_args = ssh_args + [ + "ssh", + "-o", + "UserKnownHostsFile=/dev/null", + "-o", + "StrictHostKeyChecking=no", + f"{user}@{host}", + ] + cmd = nix_shell_args + ["torify"] + password_args + _ssh_args + subprocess.run(cmd) + + +def qrcode_scan(pictureFile: str) -> dict: + subprocess.Popen( + [ + "nix", + "shell", + "nixpkgs#zbar", + "-c", + "zbarimg", + "--quiet", + "--raw", + pictureFile, + ], + stdout=subprocess.PIPE, + ).stdout.read() + + +def main(args: argparse.Namespace) -> None: + if args.json: + with open(args.json) as file: + ssh_data = json.load(file) + ssh(host=ssh_data["address"], password=ssh_data["password"]) + elif args.png: + ssh_data = json.loads(qrcode_scan(args.png)) + ssh(host=ssh_data["address"], password=ssh_data["password"]) + + +def register_parser(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "-j", + "--json", + help="specify the json file for ssh data (generated by starting the clan installer", + ) + parser.add_argument( + "-P", + "--png", + help="specify the json file for ssh data as the qrcode image (generated by starting the clan installer", + ) + # TODO pass all args we don't parse into ssh_args, currently it fails if arg starts with - + parser.add_argument("ssh_args", nargs="*", default=[]) + parser.set_defaults(func=main) diff --git a/pkgs/clan-cli/tests/test_clan_ssh.py b/pkgs/clan-cli/tests/test_clan_ssh.py new file mode 100644 index 00000000..99f91483 --- /dev/null +++ b/pkgs/clan-cli/tests/test_clan_ssh.py @@ -0,0 +1,57 @@ +import argparse +import json +import tempfile +from typing import Union + +import pytest_subprocess.fake_process +from pytest_subprocess import utils + +import clan_cli.ssh + + +# using fp fixture from pytest-subprocess +def test_ssh_no_pass(fp: pytest_subprocess.fake_process.FakeProcess) -> None: + host = "somehost" + user = "user" + cmd: list[Union[str, utils.Any]] = [ + "torify", + "ssh", + "-o", + "UserKnownHostsFile=/dev/null", + "-o", + "StrictHostKeyChecking=no", + f"{user}@{host}", + ] + fp.register(cmd) + clan_cli.ssh.ssh( + host=host, + user=user, + ) + assert fp.call_count(cmd) == 1 + + +# using fp fixture from pytest-subprocess +def test_ssh_json(fp: pytest_subprocess.fake_process.FakeProcess) -> None: + with tempfile.NamedTemporaryFile(mode="w+") as file: + json.dump({"password": "XXX", "address": "somehost"}, file) + cmd: list[Union[str, utils.Any]] = [ + "nix", + "shell", + "nixpkgs#sshpass", + "-c", + "torify", + "sshpass", + "-p", + "XXX", + "ssh", + "-o", + "UserKnownHostsFile=/dev/null", + "-o", + "StrictHostKeyChecking=no", + "root@somehost", + ] + fp.register(cmd) + file.seek(0) # write file and go to the beginning + args = argparse.Namespace(json=file.name, ssh_args=[]) + clan_cli.ssh.main(args) + assert fp.call_count(cmd) == 1