From 1f244e38596ea8134b4e19739fb76265b7df4a47 Mon Sep 17 00:00:00 2001
From: Ben Sima <ben@bsima.me>
Date: Mon, 28 Nov 2022 12:27:05 -0500
Subject: init working

---
 .envrc           |   1 +
 .gitignore       |   2 +
 nix-build-omatic | 128 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 shell.nix        |   9 ++++
 4 files changed, 140 insertions(+)
 create mode 100644 .envrc
 create mode 100644 .gitignore
 create mode 100755 nix-build-omatic
 create mode 100644 shell.nix

diff --git a/.envrc b/.envrc
new file mode 100644
index 0000000..1d953f4
--- /dev/null
+++ b/.envrc
@@ -0,0 +1 @@
+use nix
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..2bbdbfe
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+.direnv
+result
diff --git a/nix-build-omatic b/nix-build-omatic
new file mode 100755
index 0000000..ce8a7bb
--- /dev/null
+++ b/nix-build-omatic
@@ -0,0 +1,128 @@
+#!/usr/bin/env python3
+"""Provision a large virtual machine on Digital Ocean and run a nix remote
+build, then destroy the VM.
+
+I tested this like so:
+
+   $ ./nix-build-omatic "'<nixpkgs>'" -A ocaml
+
+The special quoting is necessary until I setup an argparse CLI, currently it
+just passes all args through to the nix-build call. But it works!
+
+"""
+
+import json
+import os
+import sys
+import subprocess
+import time
+
+def noti(msg):
+    print(f"nix-build-omatic: {msg}")
+
+class Builder:
+    def __init__(self):
+        # TODO: generate a temporary privkey, and use that instead
+        self.keys = [k["id"] for k in json.loads(subprocess.check_output(["doctl", "compute", "ssh-key", "list", "--output", "json"]).decode('utf-8'))]
+        self.privkey_path = os.path.expanduser("~/.ssh/id_rsa")
+        self.max_jobs = 32
+        self.ip = None
+
+    def __enter__(self):
+        "Startup a droplet."
+        noti("preparing droplet...")
+        res = self.droplet(
+            "create",
+            # this is my custom nixos image, I had to download from nixos.org and then
+            # upload it to DO. There's gotta be a better way to do this..
+            "--image", "121034270",
+            "--size", "c2-32vcpu-64gb", # TODO: set this in argparse, match '32vcpu' with 'max_jobs'
+            "--region", "nyc3",
+            "--tag-name", "nix-build-omatic",
+            "--ssh-keys", ",".join([str(k) for k in self.keys]),
+            "nix-build-omatic"
+        )
+        self.id = str(res["id"])
+        self.active = False
+        while not self.active:
+            res = self.droplet("get", self.id)
+            if res["status"] == "active":
+                self.active = True
+                noti("droplet ready!")
+            else:
+                continue
+        while self.ip is None:
+            res = self.droplet("get", self.id)
+            self.ip = [n["ip_address"] for n in res["networks"]["v4"] if n["type"] == "public" ][0]
+        self.url = f"ssh://root@{self.ip}"
+        self.do_stupid_handshake()
+        return self
+
+    def __exit__(self, type, value, traceback):
+        "Destroy the droplet."
+        self.droplet("delete", self.id, "--force", parse=False)
+        noti("droplet destroyed")
+
+    def droplet(self, *args, parse=True):
+        "Run a `doctl compute droplet *args` command"
+        proc = subprocess.run(
+            ["doctl", "compute", "droplet", *args, "--output", "json"],
+            capture_output=True
+        )
+        stderr = proc.stderr.decode('utf-8')
+        stdout = proc.stdout.decode('utf-8')
+        if proc.returncode > 0:
+            noti("stderr:", stderr)
+            noti("stdout:", stdout)
+        if parse:
+            return json.loads(stdout)[0]
+        else:
+            return stdout
+
+    def do_stupid_handshake(self):
+        """
+        this whole thing is doing a key handshake, which shouldn't be necessary!
+        i should just generate the keys on my end and set them via the API! or
+        retrieve them via the API! i'm already authenticated! wtf digital ocean
+        """
+        noti("exchanging keys...")
+        keyscan = None
+        while keyscan is None:
+            cmd = f"ssh-keyscan -t ed25519 {self.ip}"
+            proc = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+            stderr = proc.stderr.decode('utf-8')
+            stdout = proc.stdout.decode('utf-8')
+            if proc.returncode > 0:
+                continue
+            elif stdout == "":
+                continue
+            else:
+                keyscan = stdout
+        with open(os.path.expanduser("~/.ssh/known_hosts"), 'a') as known_hosts:
+            known_hosts.write(keyscan)
+        self.host_pubkey = subprocess.check_output(f"ssh root@{self.ip} cat /etc/ssh/ssh_host_ed25519_key.pub | base64 -w0", shell=True).decode('utf-8')
+
+
+if __name__ == '__main__':
+    # TODO: use argparse, properly quote '<nixpkgs>' and so on
+    args = sys.argv[1:]
+    start_time = time.time()
+
+    with Builder() as b:
+        cmd = " ".join([
+            "nix-build",
+            *args,
+            # TODO: obv get rid of this when I'm not testing...
+            "--option", "substitute", "false",
+            # TODO: make this --builders line structured data somehow
+            "--builders", f"'{b.url} x86_64-linux {b.privkey_path} {b.max_jobs} - - - {b.host_pubkey}'"
+        ])
+        noti(f"running: {cmd}")
+        subprocess.check_call(cmd, shell=True)
+
+    duration = time.time() - start_time
+    duration_human = time.strftime('%H:%M:%S', duration)
+    # the 32 vCPU machine costs $1/hr, so convert that to cost per second
+    cost = (1/60/60)*duration
+    print(f"duration: {duration_human}")
+    print(f"estimated cost: {cost:0.2f} USD")
diff --git a/shell.nix b/shell.nix
new file mode 100644
index 0000000..480dc89
--- /dev/null
+++ b/shell.nix
@@ -0,0 +1,9 @@
+with import <nixpkgs> {};
+
+mkShell {
+  name = "nix-build-ephemeral dev";
+  buildInputs = with pkgs; [
+    python3
+    doctl
+  ];
+}
-- 
cgit v1.2.3