#!/usr/bin/env python3
# SPDX-License-Identifier: GPL-3.0-or-later
import hashlib,ipaddress,json,os,socket,struct,time
from pathlib import Path
ELF_SIZE=109904;ELF_SHA="b41763f9261b838be2f83b0ee1c7b9bb5e07ced856e5638e3357de37e4405394";PACKAGE_SIZE=18153472;PACKAGE_SHA="dbcdd4dbc6303fc7a94aa0e8bb3e2c7de1d8b5770c6a30cf1ce50bc6e373aa7e";FRAME=struct.Struct("<8sIIiiQQQQq");PORT=9021;MAX_WIRE=PACKAGE_SIZE+128
FIELDS={"active","run_id","target","port","artifact_size","artifact_sha256","output_path","receipt_path","not_before","not_after","one_connection","one_transfer","one_execution","result_receive","exact_package_read","possible_atime_effect_acknowledged","device_write","usb_read","backup_read","installation","autoload","retry","reconnect"}
class Error(RuntimeError):pass
def load(p):return json.loads(p.read_text())
def validate(a,b,now):
 if set(a)!=FIELDS or set(b)!=FIELDS or a!=b or a["active"] is not True:raise Error("inactive/mismatch")
 if not a["not_before"]<=now<=a["not_after"] or a["port"]!=PORT or a["artifact_size"]!=ELF_SIZE or a["artifact_sha256"]!=ELF_SHA:raise Error("window/identity")
 if any(a[k] is not True for k in ("one_connection","one_transfer","one_execution","result_receive","exact_package_read","possible_atime_effect_acknowledged")) or any(a[k] is not False for k in ("device_write","usb_read","backup_read","installation","autoload","retry","reconnect")):raise Error("authority")
 ipaddress.IPv4Address(a["target"])
 for k in ("output_path","receipt_path"):
  p=Path(a[k]);
  if not p.is_absolute() or p.exists():raise Error("exclusive path")
 return a
def put(p,v):
 with p.open("x",encoding="utf-8",newline="\n") as f:json.dump(v,f,sort_keys=True);f.write("\n");f.flush();os.fsync(f.fileno())
def parse_frame(raw):
 v=FRAME.unpack(raw)
 if v[0]!=b"CHP10DW1" or v[1]!=1:raise Error("frame identity")
 return v
def run(elf,activation,approval,factory=socket.socket,clock=time.monotonic,wall=time.time):
 p=validate(load(activation),load(approval),wall());raw=elf.read_bytes()
 if len(raw)!=ELF_SIZE or hashlib.sha256(raw).hexdigest()!=ELF_SHA:raise Error("artifact")
 put(Path(p["receipt_path"]),{"consumed_before_socket":True,"run_id":p["run_id"],"artifact_sha256":ELF_SHA});s=None;out=None;header=bytearray();terminal=bytearray();written=0;digest=hashlib.sha256();deadline=clock()+45
 try:
  out=Path(p["output_path"]).open("xb");s=factory(socket.AF_INET,socket.SOCK_STREAM);s.settimeout(3);s.connect((p["target"],PORT));s.sendall(raw);s.shutdown(socket.SHUT_WR)
  while len(header)<64:
   s.settimeout(max(.001,deadline-clock()));c=s.recv(64-len(header))
   if not c:raise Error("header eof")
   header.extend(c)
  h=parse_frame(bytes(header));
  if h[2]!=1 or h[3] or h[4] or h[5]!=PACKAGE_SIZE or h[6]:raise Error("header")
  while written<PACKAGE_SIZE:
   if clock()>=deadline:raise Error("deadline")
   s.settimeout(max(.001,deadline-clock()));c=s.recv(min(65536,PACKAGE_SIZE-written))
   if not c:raise Error("payload eof")
   out.write(c);digest.update(c);written+=len(c)
  while len(terminal)<64:
   s.settimeout(max(.001,deadline-clock()));c=s.recv(64-len(terminal))
   if not c:raise Error("terminal eof")
   terminal.extend(c)
  if s.recv(1):raise Error("trailing data")
  t=parse_frame(bytes(terminal));
  if t[2]!=2 or t[3] or t[4] or t[5]!=PACKAGE_SIZE or t[6]!=PACKAGE_SIZE or t[7:10]!=h[7:10]:raise Error("terminal")
  out.flush();os.fsync(out.fileno());out.close();out=None
  if digest.hexdigest()!=PACKAGE_SHA:raise Error("package digest mismatch")
  return {"size":written,"sha256":digest.hexdigest()}
 finally:
  if out is not None:out.close()
  if s is not None:s.close()