rewrite remote_as_root in python

diff --git a/remote_as_root.sh b/remote_as_root.sh
deleted file mode 100755 (executable)
index 5f441c6..0000000
--- a/remote_as_root.sh
+++ /dev/null
@@ -1,37 +0,0 @@
-#!/usr/bin/env bash
-
-# this works if the remote has pam_ssh_agent_auth with the local id_ed25519.pub set up
-
-script="$(realpath -e "${0}")"
-script_dir="${script%/*}"
-
-get_ssh_agent_socket() {
-    local tmpdir="${TMPDIR:-/tmp}"
-    local former_nullglob_set=0 found=0 output
-    shopt -q nullglob && former_nullglob_set=1
-    shopt -s nullglob
-    for SSH_AUTH_SOCK in "${tmpdir}"/ssh-*/agent.*; do
-        output="$(ssh-add -l 2>&1)"
-        [[ "${output}" == "Error connecting to agent: Connection refused" ]] && continue
-        export SSH_AUTH_SOCK
-        found=1
-        break
-    done
-    (( found == 0 )) && . <(ssh-agent)
-    (( former_nullglob_set )) || shopt -u nullglob
-    return 0
-}
-
-add_ssh_key() {
-    local privkey="${script_dir}/${1}"
-    [[ ! -f "${privkey}" ]] && ssh-keygen -t "${1}" -b 521 -f "${privkey}"
-    grep -qE "^$(ssh-keygen -lf "${privkey}"|sed -r 's/([][{}|()+?])/\\\1/g')\$" \
-        <(ssh-add -l) || ssh-add "${privkey}"
-}
-
-get_ssh_agent_socket
-add_ssh_key id_ed25519
-
-script_b64="$(gzip -9<"${2}"|base64 -w0)"
-launcher="$(printf '%q' '. <(printf "%s" "'"${script_b64}"'"|base64 -d|gunzip)')"
-ssh -o ForwardX11=no -At "${1}" sudo bash -c "${launcher}"

diff --git a/remote_run.py b/remote_run.py
new file mode 100755 (executable)
index 0000000..8065fa4
--- /dev/null
+++ b/remote_run.py
@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+
+import os
+import re
+import sys
+from argparse import ArgumentError
+from itertools import chain
+from pathlib import Path
+from socketserver import BaseRequestHandler, TCPServer
+from subprocess import DEVNULL, PIPE, Popen, run
+
+PORT = "7777"
+REMOTE_HOST = "localhost"
+REMOTE_PORT = "7777"
+AUTH_SOCK_DIR_PATTERN = re.compile(r"/ssh-[^/]+$")
+AUTH_SOCK_PATTERN = re.compile(
+    rf"{AUTH_SOCK_DIR_PATTERN.pattern.rstrip('$')}/agent.[0-9]+$"
+)
+
+
+def check_ssh_auth_sock(**kwargs):
+    completed = run(
+        ["ssh-add", "-l"],
+        stdout=PIPE,
+        stderr=DEVNULL,
+        universal_newlines=True,
+        **kwargs,
+    )
+    return (
+        completed.returncode == 0
+        or completed.stdout.rstrip() == "The agent has no identities."
+    )
+
+
+def get_ssh_auth_sock():
+    if "SSH_AUTH_SOCK" in os.environ and check_ssh_auth_sock():
+        return os.environ["SSH_AUTH_SOCK"]
+    for q in chain.from_iterable(
+        p.iterdir()
+        for p in Path(os.environ.get("TMPDIR") or "/tmp").iterdir()
+        if p.is_dir() and AUTH_SOCK_DIR_PATTERN.search(str(p))
+    ):
+        q_str = str(q)
+        if (
+            q.is_socket()
+            and AUTH_SOCK_PATTERN.search(q_str)
+            and check_ssh_auth_sock(env={**os.environ, "SSH_AUTH_SOCK": q_str})
+        ):
+            return q_str
+    print("starting new ssh-agent")
+    return run(
+        ["bash", "-c", '. <(ssh-agent); echo "${SSH_AUTH_SOCK}"'],
+        stdout=PIPE,
+        universal_newlines=True,
+    ).stdout
+
+
+def create_tcp_server(server_address, data):
+    if not isinstance(data, bytes):
+        data = data.encode()
+
+    class RequestHandler(BaseRequestHandler):
+        def handle(self):
+            self.request.sendall(data)
+
+    class ReusableTCPServer(TCPServer):
+        allow_reuse_address = True
+
+    return ReusableTCPServer(server_address, RequestHandler)
+
+
+def usage():
+    print(f"{sys.argv[0]} [-h|--help] [--sudo] ssh-args script")
+    print("\t-h, --help  display this help message")
+    print("\t--sudo      run remote script as root")
+    print("\tssh-args    arguments passed to ssh; include a hostname here")
+    print("\tscript      script to run remotely")
+
+
+def main():
+    args = sys.argv[1:-1]
+    if len(args) < 1:
+        raise ArgumentError(None, "No host specified")
+    if "--help" in args or "-h" in args:
+        usage()
+        exit(0)
+    if "--sudo" in args:
+        sudo = "sudo "
+        args.remove("--sudo")
+    else:
+        sudo = ""
+    environ = {**os.environ, "SSH_AUTH_SOCK": get_ssh_auth_sock()}
+    run(["ssh-add", "id_ed25519"], stderr=DEVNULL, env=environ)
+    with open(sys.argv[-1], "rb") as fh:
+        with create_tcp_server(("127.0.0.1", int(PORT)), fh.read()) as server:
+            sp = Popen(
+                [
+                    "ssh",
+                    "-R",
+                    ":".join((PORT, REMOTE_HOST, REMOTE_PORT)),
+                    "-o",
+                    "ForwardX11=no",
+                    "-At",
+                    *args,
+                    f"{sudo}bash -c '. <(cat </dev/tcp/{REMOTE_HOST}/{REMOTE_PORT})'",
+                ],
+                env=environ,
+            )
+            server.handle_request()
+            sp.communicate()
+
+
+if __name__ == "__main__":
+    main()