From 95a64e4705b45ce9742595b8c570d5c6d67aae7c Mon Sep 17 00:00:00 2001
From: Drew DeVault <sir@cmpwn.com>
Date: Thu, 30 Jul 2020 16:59:31 -0400
Subject: [PATCH] Add "master-shell", supports submitting via SSH

---
 buildsrht-keys                  | 30 +++++++++++++-----
 master-shell                    | 56 +++++++++++++++++++++++++++++++++
 buildsrht-shell => runner-shell | 13 +-------
 setup.py                        |  3 +-
 4 files changed, 81 insertions(+), 21 deletions(-)
 create mode 100755 master-shell
 rename buildsrht-shell => runner-shell (88%)

diff --git a/buildsrht-keys b/buildsrht-keys
index 6692d41..0931632 100755
--- a/buildsrht-keys
+++ b/buildsrht-keys
@@ -1,16 +1,30 @@
 #!/usr/bin/env python3
-# We just let everyone in at this stage, authentication is done later on.
-from srht.config import cfg
-import sys
+from srht.config import cfg, get_origin
 import os
+import requests
+import sys
 
 key_type = sys.argv[3]
 b64key = sys.argv[4]
 
-default_shell = os.path.join(os.path.dirname(sys.argv[0]), "buildsrht-shell")
+def fail(reason):
+    owner = cfg("sr.ht", "owner-name")
+    email = cfg("sr.ht", "owner-email")
+    print(reason)
+    print(f"Please reach out to {owner} <{email}> for support.")
+    sys.exit(1)
+
+meta_origin = get_origin("meta.sr.ht")
+r = requests.get(f"{meta_origin}/api/ssh-key/{b64key}")
+if r.status_code == 200:
+    username = r.json()["owner"]["name"]
+elif r.status_code == 404:
+    fail("We don't recognize your SSH key. Make sure you've added it to " +
+        f"your account.\n{get_origin('meta.sr.ht', external=True)}/keys")
+else:
+    fail("Temporary authentication failure. Try again later.")
+
+default_shell = os.path.join(os.path.dirname(sys.argv[0]), "master-shell")
 shell = cfg("builds.sr.ht", "shell", default=default_shell)
-keys = ("command=\"{} '{}'\",".format(shell, b64key) +
-    "no-port-forwarding,no-X11-forwarding,no-agent-forwarding" +
-    " {} {} somebody".format(key_type, b64key) + "\n")
-print(keys)
+print(f"restrict,pty,command=\"{shell} '{username}'\" {key_type} {b64key} somebody\n")
 sys.exit(0)
diff --git a/master-shell b/master-shell
new file mode 100755
index 0000000..9ad4bcc
--- /dev/null
+++ b/master-shell
@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+from buildsrht.manifest import Manifest
+from buildsrht.runner import queue_build
+from buildsrht.types import Job, Task, User
+from getopt import getopt, GetoptError
+from srht.config import cfg, get_origin
+from srht.database import DbSession
+import os
+import shlex
+import sys
+import yaml
+
+db = DbSession(cfg("builds.sr.ht", "connection-string"))
+db.init()
+
+def fail(reason):
+    owner = cfg("sr.ht", "owner-name")
+    email = cfg("sr.ht", "owner-email")
+    print(reason)
+    print(f"Please reach out to {owner} <{email}> for support.")
+    sys.exit(1)
+
+username = sys.argv[1]
+user = User.query.filter(User.username == username).one_or_none()
+if not user:
+    fail(f"Unknown user {username}")
+cmd = os.environ.get("SSH_ORIGINAL_COMMAND") or ""
+cmd = shlex.split(cmd)
+
+if cmd[0] == "submit":
+    try:
+        opts, args = getopt(cmd[1:], "n:")
+    except GetoptError as ex:
+        fail(str(ex))
+
+    if os.isatty(sys.stdin.fileno()):
+        print("Enter build manifest:")
+    _manifest = sys.stdin.read()
+    try:
+        manifest = Manifest(yaml.safe_load(_manifest))
+    except Exception as ex:
+        fail(str(ex))
+    job = Job(user, _manifest)
+    job.image = manifest.image
+    job.note = [y for x, y in opts if x == "-n"] or None
+    db.session.add(job)
+    db.session.flush()
+    for task in manifest.tasks:
+        t = Task(job, task.name)
+        db.session.add(t)
+        db.session.flush() # assigns IDs for ordering purposes
+    queue_build(job, manifest) # commits the session
+    url = f"{get_origin('builds.sr.ht', external=True)}/~{username}/job/{job.id}"
+    print(url)
+else:
+    fail(f"Unknown command {cmd[0]}")
diff --git a/buildsrht-shell b/runner-shell
similarity index 88%
rename from buildsrht-shell
rename to runner-shell
index 2d910c6..9178b1f 100755
--- a/buildsrht-shell
+++ b/runner-shell
@@ -19,8 +19,6 @@ def fail(reason):
     print(f"Please reach out to {owner} <{email}> for support.")
     sys.exit(1)
 
-b64key = sys.argv[1]
-
 cmd = os.environ.get("SSH_ORIGINAL_COMMAND") or ""
 cmd = shlex.split(cmd)
 if len(cmd) < 2:
@@ -43,16 +41,7 @@ info = get_info(job_id)
 if not info:
     fail("No such job found.")
 
-meta_origin = get_origin("meta.sr.ht")
-r = requests.get(f"{meta_origin}/api/ssh-key/{b64key}")
-if r.status_code == 200:
-    username = r.json()["owner"]["name"]
-elif r.status_code == 404:
-    fail("We don't recognize your SSH key. Make sure you've added it to " +
-        f"your account.\n{get_origin('meta.sr.ht', external=True)}/keys")
-else:
-    fail("Temporary authentication failure. Try again later.")
-
+username = sys.argv[1]
 if username != info["username"]:
     fail("You are not permitted to connect to this job.")
 
diff --git a/setup.py b/setup.py
index cdd6d2a..a763302 100755
--- a/setup.py
+++ b/setup.py
@@ -50,6 +50,7 @@ setup(
       'buildsrht-initdb',
       'buildsrht-keys',
       'buildsrht-migrate',
-      'buildsrht-shell',
+      'master-shell',
+      'runner-shell',
   ]
 )