~sircmpwn/builds.sr.ht

builds.sr.ht/runner-shell -rwxr-xr-x 3.4 KiB
570f7196Francesco Gazzetta Remove end-of-life NixOS 20.09 15 days ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
#!/usr/bin/env python3
from buildsrht.manifest import Manifest
from datetime import datetime
from humanize import naturaltime
from srht.config import cfg, get_origin
from srht.redis import redis
import os
import requests
import shlex
import subprocess
import sys
import time
import yaml

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)

cmd = os.environ.get("SSH_ORIGINAL_COMMAND") or ""
cmd = shlex.split(cmd)
if len(cmd) < 2:
    fail("Usage: ssh ... connect <job ID>")
op = cmd[0]
if op not in ["connect", "tail"]:
    fail("Usage: ssh ... connect <job ID>")
job_id = int(cmd[1])
cmd = cmd[2:]

bind_address = cfg("builds.sr.ht::worker", "bind-address", "0.0.0.0:8080")

def get_info(job_id):
    r = requests.get(f"http://{bind_address}/job/{job_id}/info")
    if r.status_code != 200:
        return None
    return r.json()

info = get_info(job_id)
if not info:
    fail("No such job found.")

username = sys.argv[1]
if username != info["username"]:
    fail("You are not permitted to connect to this job.")

if len(cmd) == 0:
    url = f"{get_origin('builds.sr.ht', external=True)}/~{username}/job/{job_id}"
    print(f"Connected to build job #{job_id} ({info['status']}): {url}")
deadline = datetime.utcfromtimestamp(info["deadline"])

manifest = Manifest(yaml.safe_load(info["manifest"]))

def connect(job_id, info):
    """Opens a shell on the build VM"""
    limit = naturaltime(datetime.utcnow() - deadline)
    if len(cmd) == 0:
        print(f"Your VM will be terminated {limit}, or when you log out.")
        print()
    requests.post(f"http://{bind_address}/job/{job_id}/claim")
    sys.stdout.flush()
    sys.stderr.flush()
    try:
        tty = os.open("/dev/tty", os.O_RDWR)
        os.dup2(0, tty)
    except:
        pass # non-interactive
    redis.incr(f"builds.sr.ht-shell-{job_id}")
    subprocess.call([
        "ssh", "-qt",
        "-p", str(info["port"]),
        "-o", "UserKnownHostsFile=/dev/null",
        "-o", "StrictHostKeyChecking=no",
        "-o", "LogLevel=quiet",
        "build@localhost",
    ] + cmd)
    n = redis.decr(f"builds.sr.ht-shell-{job_id}")
    if n == 0:
        requests.post(f"http://{bind_address}/job/{job_id}/terminate")

def tail(job_id, info):
    """Tails the build logs to stdout"""
    logs = os.path.join(cfg("builds.sr.ht::worker", "buildlogs"), str(job_id))
    p = subprocess.Popen(["tail", "-f", os.path.join(logs, "log")])
    tasks = set()
    procs = [p]
    # holy bejeezus this is hacky
    while True:
        for task in manifest.tasks:
            if task.name in tasks:
                continue
            path = os.path.join(logs, task.name, "log")
            if os.path.exists(path):
                procs.append(subprocess.Popen(
                    f"tail -f {shlex.quote(path)} | " +
                    "awk '{ print \"[" + shlex.quote(task.name) + "] \" $0 }'",
                    shell=True))
                tasks.update({ task.name })
        info = get_info(job_id)
        if not info:
            break
        if info["task"] == info["tasks"]:
            for p in procs:
                p.kill()
            break
        time.sleep(3)

if op == "connect":
    if info["task"] != info["tasks"] and info["status"] == "running":
        tail(job_id, info)
    connect(job_id, info)
elif op == "tail":
    tail(job_id, info)