git.mar77i.info Git - admin/blob - remote_run.py

   1 #!/usr/bin/env python3
   2
   3 import os
   4 import re
   5 import sys
   6 from argparse import ArgumentError
   7 from itertools import chain
   8 from pathlib import Path
   9 from socketserver import BaseRequestHandler, TCPServer
  10 from subprocess import DEVNULL, PIPE, Popen, check_output, run
  11
  12 PORT = "7777"
  13 REMOTE_HOST = "localhost"
  14 REMOTE_PORT = "7777"
  15 AUTH_SOCK_DIR_PATTERN = re.compile(r"ssh-")
  16 AUTH_SOCK_PATTERN = re.compile(f"agent.[0-9]+$")
  17
  18
  19 class ReusableTCPServer(TCPServer):
  20     allow_reuse_address = True
  21
  22
  23 def check_ssh_auth_sock(**kwargs):
  24     completed = run(
  25         ["ssh-add", "-l"],
  26         stdout=PIPE,
  27         stderr=DEVNULL,
  28         universal_newlines=True,
  29         **kwargs,
  30     )
  31     return (
  32         completed.returncode == 0
  33         or completed.stdout.rstrip() == "The agent has no identities."
  34     )
  35
  36
  37 def get_ssh_auth_sock():
  38     if "SSH_AUTH_SOCK" in os.environ and check_ssh_auth_sock():
  39         return os.environ["SSH_AUTH_SOCK"]
  40     for q in chain.from_iterable(
  41         p.iterdir()
  42         for p in Path(os.environ.get("TMPDIR") or "/tmp").iterdir()
  43         if p.is_dir() and AUTH_SOCK_DIR_PATTERN.match(p.name)
  44     ):
  45         q_str = str(q)
  46         if (
  47             q.is_socket()
  48             and AUTH_SOCK_PATTERN.fullmatch(q.name)
  49             and check_ssh_auth_sock(env={**os.environ, "SSH_AUTH_SOCK": q_str})
  50         ):
  51             return q_str
  52     print("starting new ssh-agent")
  53     return check_output(
  54         ["bash", "-c", '. <(ssh-agent); echo "${SSH_AUTH_SOCK}"'],
  55         universal_newlines=True,
  56     )
  57
  58
  59 def create_tcp_server(server_address, data):
  60     if not isinstance(data, bytes):
  61         data = data.encode()
  62
  63     class RequestHandler(BaseRequestHandler):
  64         def handle(self):
  65             self.request.sendall(data)
  66
  67     return ReusableTCPServer(server_address, RequestHandler)
  68
  69
  70 def usage():
  71     print(f"{sys.argv[0]} [-h|--help] [--sudo] ssh-args script")
  72     print("\t-h, --help  display this help message")
  73     print("\t--sudo      run remote script as root")
  74     print("\tssh-args    arguments passed to ssh; include a hostname here")
  75     print("\tscript      script to run remotely")
  76
  77
  78 def main():
  79     args = sys.argv[1:-1]
  80     if len(args) < 1:
  81         raise ArgumentError(None, "No host specified")
  82     if "--help" in args or "-h" in args:
  83         usage()
  84         exit(0)
  85     if "--sudo" in args:
  86         sudo = "sudo "
  87         args.remove("--sudo")
  88     else:
  89         sudo = ""
  90     environ = {**os.environ, "SSH_AUTH_SOCK": get_ssh_auth_sock()}
  91     run(["ssh-add", "id_ed25519"], stderr=DEVNULL, env=environ)
  92     with open(sys.argv[-1], "rb") as fh:
  93         with create_tcp_server(("127.0.0.1", int(PORT)), fh.read()) as server:
  94             sp = Popen(
  95                 [
  96                     "ssh",
  97                     "-R",
  98                     ":".join((PORT, REMOTE_HOST, REMOTE_PORT)),
  99                     "-o",
 100                     "ForwardX11=no",
 101                     "-At",
 102                     *args,
 103                     f"{sudo}bash -c '. <(cat </dev/tcp/{REMOTE_HOST}/{REMOTE_PORT})'",
 104                 ],
 105                 env=environ,
 106             )
 107             server.handle_request()
 108             sp.communicate()
 109
 110
 111 if __name__ == "__main__":
 112     main()