115 lines
3.4 KiB
Python
Executable File
115 lines
3.4 KiB
Python
Executable File
#!/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)
|