main.py

import os
import random
import re
import secrets
import socket
import threading
import time

import colorama
import requests

from subshell import Shell

colorama.init()

from config import *


class TestData:
  def __init__(self):
    self.pin = random.randint(0, 999999)
    self.token = secrets.token_hex(8)
    self.component_count = 2
    self.component_ids = [None] * self.component_count
    for i in range(self.component_count):
      rand_comp = secrets.randbits(32)
      while rand_comp & 0xff < 0x8 or rand_comp & 0xff > 0x78:
        rand_comp = secrets.randbits(32)
      self.component_ids[i] = "0x{:08x}".format(rand_comp)
    self.boot_message_ap = f"boot_msg ap ({secrets.token_hex(8)})"
    self.component_dat = [(f"boot_msg comp{_} ({secrets.token_hex(8)})",  # boot msg
                      secrets.token_hex(4), # attestation location
                      secrets.token_hex(8), # attestation cust
                      f"{random.randint(1, 12)}/{random.randint(1, 31)}/{random.randint(1950, 2024)}") # attestation date
                      for _ in range(self.component_count)]

class CommitInfo:
  def __init__(self, hash:str, author:str, name:str, run_id: str):
    self.hash = hash
    self.name = name
    self.author = author
    self.run_id = run_id
    self.out_path = hash[:7]
  def to_json(self):
    return {
      "hash": self.hash,
      "name": self.name,
      "author": self.author,
      "runId": self.run_id,
    }

class ActionResult:
  def __init__(self, conn: socket.socket, info: CommitInfo, status: str, time: int, test_data: TestData = None):
    self.conn = conn
    self.info = info
    self.data = test_data
    self.status = status
    self.time = time
  def to_json(self):
    return {
      "result": self.status,
      "actionStart": round(self.time),
      "commit": self.info.to_json()
    }

class TestServerStatus:
  def __init__(self, locked: bool, result: ActionResult):
    self.locked = locked
    self.result = result
  def is_avail(self):
    return not self.locked or not self.result or self.result.status != "TESTING"

build_queue: list[ActionResult] = []
upload_queue: list[ActionResult] = []
active_build: ActionResult = None
upload_status: dict[str, TestServerStatus] = {}

build_shell = Shell()
last_hash = "release"

def push_webhook(update_type: str = "QUEUE", update_state: ActionResult = None):
  requests.post(WEBHOOK_IP, json={
    "update": {
      "type": update_type,
      "state": update_state.to_json() if update_state else None,
    },
    "build": {
      "active": active_build.to_json() if active_build else None,
      "queue": [action.to_json() for action in build_queue],
    },
    "test": {
      "activeTests": [
        {"ip": ip, "locked": stat.locked, "active": stat.result.to_json() if stat.result else None}
        for ip, stat in upload_status.items()],
      "queue": [action.to_json() for action in upload_queue],
    }
  }, headers={
    "content-type": "application/json"
  })


def upload(req: ActionResult, ip: str):
  req.conn.send((colorama.Fore.BLUE + f"Uploading commit {req.info.hash}..." + colorama.Fore.RESET + "\n").encode())
  print(f"Uploading commit {req.info.hash} to {ip}...")
  sh = Shell()
  sh.output_conn = req.conn
  try:
    sh.run_cmd(f'rsync --rsh="ssh -o StrictHostKeyChecking=accept-new" -av --progress --rsync-path="rm -rf ~/CI/rpi/upload/{req.info.out_path}; mkdir -p ~/CI/rpi/upload/{req.info.out_path} && rsync" 2024-ectf-secure-example/build/{req.info.out_path}/{{*.img,test_data}} {ip}:~/CI/rpi/upload/{req.info.out_path}')
  except Exception:
    print("Failed to upload!")
    req.conn.send((colorama.Fore.RED + f"Failed to upload!" + colorama.Fore.RESET + "\n").encode())
    req.conn.send(b"%*&1\n")
    req.conn.close()
    sh.close()
    return

  sh.run_cmd(f'rm -rf 2024-ectf-secure-example/build/{req.info.out_path}')
  req.conn.send((colorama.Fore.BLUE + f"Uploaded! Running tests..." + colorama.Fore.RESET + "\n").encode())
  print(f"Running tests for {req.info.hash}...")
  try:
    sh.run_cmd(f"ssh {ip} \"bash -l -c \\\"cd CI/rpi/;chmod +x run-tests.sh; nix-shell --run './run-tests.sh {req.info.out_path}' \\\" \"")
  except Exception:
    req.conn.send((colorama.Fore.RED + f"Tests failed!" + colorama.Fore.RESET + "\n").encode())
    print(f"Tests failed for {req.info.hash}")
    req.conn.send(b"%*&1\n")
    req.status = "FAILED"
    req.time = time.time()
    push_webhook("TEST", req)
    return
  print(f"Tests OK for {req.info.hash}")
  req.conn.send(b"%*&0\n")
  req.conn.close()
  req.status = "SUCCESS"
  req.time = time.time()
  push_webhook("TEST", req)
  sh.close()


def upload_loop():
  while True:
    if len(upload_queue) == 0:
      time.sleep(0.1)
      continue
    avail_ip = None
    for ip, server in upload_status.items():
      if server.is_avail():
        avail_ip = ip
        break
    if len(avail_ip) == None:
      time.sleep(0.1)
      continue
    req = upload_queue.pop(0)
    req.status = "TESTING"
    req.time = time.time()
    upload_status[avail_ip].result = req
    push_webhook()
    for other_req in upload_queue:
      other_req.conn.send(f"There are {len(upload_queue) + len(IPS) - len(ip_queue)} uploads in queue. \n".encode())
    threading.Thread(target=upload, args=(req, avail_ip), daemon=True).start()


def build_loop():
  build_shell.log = True
  try:
    try:
      build_shell.run_cmd('git clone https://github.com/Purdue-eCTF-2024/2024-ectf-secure-example.git')
    except Exception as e:
      print("Found existing repo, reusing...")
      pass
    build_shell.run_cmd('cd 2024-ectf-secure-example')
    build_shell.run_cmd('rm -rf build')
    build_shell.run_cmd('mkdir build')
    print("Finished setup")
  except Exception as e:
    print("Couldn't pull from repo, stopping...")
    os._exit(1)
  build_shell.log = False

  global active_build
  while True:
    if len(build_queue) == 0:
      time.sleep(0.1)
      continue
    try:
      req = build_queue.pop(0)
      active_build = req
      req.time = time.time()
      req.status = "BUILDING"
      push_webhook()
      for other_req in build_queue:
        other_req.conn.send(f"There are {len(build_queue) + (1 if active_build else 0)} builds in queue. \n".encode())

      i = 2
      while True:
        try:
          build_shell.run_cmd(f"test -d build/{req.info.out_path}")
        except Exception:
          break
        req.info.out_path = f"{req.info.hash[:7]}-{i}"
        i = i + 1

      if req.info.out_path != req.info.hash[:7]:
        req.conn.send((colorama.Fore.BLUE + f"Redirected build output to {req.info.out_path}" + colorama.Fore.RESET + "\n").encode())

      build_shell.output_conn = req.conn
      req.conn.send((colorama.Fore.BLUE + f"Building commit {req.info.hash}..." + colorama.Fore.RESET + "\n").encode())
      print(colorama.Fore.BLUE + f"Building commit {req.info.hash}..." + colorama.Fore.RESET)
      try:
        build_commit(req)
      except Exception as e:
        print(colorama.Fore.RED + f"Failed to build commit {req.info.hash}: {e}" + colorama.Fore.RESET)
        req.status = "FAILED"
        req.time = time.time()
        req.conn.send(b"Failed to build!\n")
        req.conn.send(b"%*&1\n")
        req.conn.close()
        build_shell.output_conn = None
        push_webhook("BUILD", req)
        continue
      print(f"Queuing upload for commit {req.info.hash}...")
      print(colorama.Fore.BLUE + f"Waiting for upload..." + colorama.Fore.RESET)
      req.conn.send(f"There are {len(upload_queue) +  len(IPS) - len(ip_queue)} uploads in queue. \n".encode())
      upload_queue.append(req)
      req.status = "PENDING"
      req.time = time.time()
      build_shell.output_conn = None
      push_webhook()
    except:
      req.status = "FAILED"
      req.time = time.time()
      req.conn.send(b"Failed to build!\n")
      req.conn.send(b"%*&1\n")
      req.conn.close()
      build_shell.output_conn = None
      push_webhook("BUILD", req)

def build_commit(req: ActionResult):
  build_shell.run_cmd(f'git reset --hard')
  build_shell.run_cmd(f'git checkout rust_hal_impl')
  build_shell.run_cmd(f'git pull')
  global last_hash
  if last_hash != req.info.hash:
    try:
      build_shell.run_cmd(f'git checkout {req.info.hash}')
    except Exception:
      req.status = "FAILED"
      req.time = time.time()
      req.conn.send(b"No commit found!\n")
      req.conn.send(b"%*&1\n")
      print(colorama.Fore.RED + f"Failed to build commit {req.info.hash}: no commit found" + colorama.Fore.RESET)
      req.conn.send(b"%*&1\n")
      req.conn.close()
      push_webhook("BUILD", req)
      return
    build_shell.run_cmd("nix-shell --run 'cargo install; poetry install'")

  build_shell.run_cmd(f"mkdir build/{req.info.out_path}")

  # generate test data
  data = TestData()
  req.data = data
  outputs = ["compa","compb"]

  build_shell.run_cmd(f"echo 'PIN=\"{data.pin}\"' >> build/{req.info.out_path}/test_data")
  build_shell.run_cmd(f"echo 'TOKEN=\"{data.token}\"' >> build/{req.info.out_path}/test_data")
  build_shell.run_cmd(f"echo 'AP_BOOT=\"{data.boot_message_ap}\"' >> build/{req.info.out_path}/test_data")
  build_shell.run_cmd(f"echo 'COMPA_ID=\"{data.component_ids[0]}\"' >> build/{req.info.out_path}/test_data")
  build_shell.run_cmd(f"echo 'COMPB_ID=\"{data.component_ids[1]}\"' >> build/{req.info.out_path}/test_data")
  build_shell.run_cmd(f"echo 'COMPA_BOOT=\"{data.component_dat[0][0]}\"' >> build/{req.info.out_path}/test_data")
  build_shell.run_cmd(f"echo 'COMPA_ATT_LOC=\"{data.component_dat[0][1]}\"' >> build/{req.info.out_path}/test_data")
  build_shell.run_cmd(f"echo 'COMPA_ATT_CUST=\"{data.component_dat[0][2]}\"' >> build/{req.info.out_path}/test_data")
  build_shell.run_cmd(f"echo 'COMPA_ATT_DATE=\"{data.component_dat[0][3]}\"' >> build/{req.info.out_path}/test_data")
  build_shell.run_cmd(f"echo 'COMPB_BOOT=\"{data.component_dat[1][0]}\"' >> build/{req.info.out_path}/test_data")
  build_shell.run_cmd(f"echo 'COMPB_ATT_LOC=\"{data.component_dat[1][1]}\"' >> build/{req.info.out_path}/test_data")
  build_shell.run_cmd(f"echo 'COMPB_ATT_CUST=\"{data.component_dat[1][2]}\"' >> build/{req.info.out_path}/test_data")
  build_shell.run_cmd(f"echo 'COMPB_ATT_DATE=\"{data.component_dat[1][3]}\"' >> build/{req.info.out_path}/test_data")

  # gen secrets if possible
  build_shell.run_cmd(f'nix-shell --run \'poetry run python ectf_tools/build_depl.py -d .\'')

  # build the thing
  print(f"Building AP...")
  build_shell.run_cmd(f'nix-shell --run \'poetry run python ectf_tools/build_ap.py -d . -on build/{req.info.out_path}/ap -p {data.pin} -t {data.token} -c {data.component_count} -ids "{", ".join(data.component_ids)}" -b "{data.boot_message_ap}"\'')
  for i in range(data.component_count):
    print(f"Building {outputs[i]}...")
    build_shell.run_cmd(f'nix-shell --run \'poetry run python ectf_tools/build_comp.py -d . -on build/{req.info.out_path}/{outputs[i]} -id {data.component_ids[i]} -b "{data.component_dat[i][0]}" -al "{data.component_dat[i][1]}" -ac "{data.component_dat[i][2]}" -ad "{data.component_dat[i][3]}"\'')
  print(colorama.Fore.BLUE + f"Built commit {req.info.hash}!" + colorama.Fore.RESET)
  last_hash = hash

def check_close(req: ActionResult):
  while True:
    try:
      check = req.conn.recv(16)
      if not check:
        req.conn.close()
        build_queue.remove(req)
        break
    except Exception:
      break
  push_webhook()
  print(f"Removed {req.info.hash[:7]} from queue")

if __name__ == "__main__":
  # setup git
  try:
    build_shell.run_cmd("gh auth status")
  except:
    print("Setting up git...")
    build_shell.run_cmd(f"echo {GH_TOKEN} | gh auth login -p https --with-token")
    build_shell.run_cmd(f"gh auth setup-git")
  # setup ssh
  try:
    build_shell.run_cmd("mkdir -p ~/.ssh")
    build_shell.run_cmd("rm ~/.ssh/config")
  except:
    pass
  for ip in IPS:
    build_shell.run_cmd(f"echo 'Host {ip.split('@')[1]}\nProxyCommand $(which cloudflared) access ssh --hostname %h\n' >> ~/.ssh/config")
    upload_status[ip] = TestServerStatus(False, None)
  push_webhook()
  print(f"Loaded {len(IPS)} pis")

  server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  server.bind(("0.0.0.0", 8888))
  server.listen()

  ip_queue = IPS

  threading.Thread(target=build_loop, daemon=True).start()
  threading.Thread(target=upload_loop, daemon=True).start()

  while True:
    try:
      conn, addr = server.accept()

      try:
        token = conn.recv(64).decode("utf-8")
        conn.send(b"Enter auth token:")
        if token != TOKEN:
          conn.close()
        conn.send(b"Enter commit details")
        hash, author, name, run_id = conn.recv(1024).decode("utf-8").split("|")
        if len(hash) > 40 or len(hash) < 7 or re.search("[^0-9a-f]", hash):
          conn.send(b"Invalid input!")
          conn.close()
          continue
      except Exception as e:
        conn.close()
        continue

      print(f"Queuing build for commit {hash}...")
      conn.send(f"There are {len(build_queue) + (1 if active_build else 0)} builds in queue. \n".encode())
      req = ActionResult(conn, CommitInfo(hash, author, name, run_id), "PENDING", time.time())
      build_queue.append(req)
      push_webhook()
      threading.Thread(target=check_close, args=(req,), daemon=True).start()

    except KeyboardInterrupt:
      break