From 3b6cbb1f063035bb5e24724827666c6b15fe0a9d Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Thu, 10 Nov 2016 20:25:15 +0000
Subject: [PATCH] artiq_devtool: implement.

---
 artiq/frontend/artiq_devtool.py | 144 ++++++++++++++++++++++++++++++++
 setup.py                        |   2 +
 2 files changed, 146 insertions(+)
 create mode 100644 artiq/frontend/artiq_devtool.py

diff --git a/artiq/frontend/artiq_devtool.py b/artiq/frontend/artiq_devtool.py
new file mode 100644
index 000000000..4c6aba681
--- /dev/null
+++ b/artiq/frontend/artiq_devtool.py
@@ -0,0 +1,144 @@
+#!/usr/bin/env python3.5
+
+# This script makes the following assumptions:
+#  * miniconda is installed remotely at ~/miniconda
+#  * misoc and artiq are installed remotely via conda
+
+import sys
+import argparse
+import subprocess
+import socket
+import select
+import threading
+import paramiko
+
+from artiq.tools import verbosity_args, init_logger, logger
+from random import Random
+
+
+def get_argparser():
+    parser = argparse.ArgumentParser(description="ARTIQ core device "
+                                                 "development tool")
+
+    verbosity_args(parser)
+
+    parser.add_argument("--host", nargs=1, metavar="HOST",
+                        type=str, default="lab.m-labs.hk",
+                        help="SSH host where the development board is located")
+    parser.add_argument("--serial", nargs=1, metavar="SERIAL",
+                        type=str, default="/dev/ttyUSB0",
+                        help="TTY device corresponding to the development board")
+    parser.add_argument("--ip", nargs=1, metavar="IP",
+                        type=str, default="kc705.lab.m-labs.hk",
+                        help="IP address corresponding to the development board")
+
+    parser.add_argument("actions", metavar="ACTION",
+                        type=str, default=[], nargs="+",
+                        help="actions to perform (sequence of: build boot connect)")
+
+    return parser
+
+
+def main():
+    args = get_argparser().parse_args()
+    init_logger(args)
+
+    ssh = None
+    def get_ssh():
+        nonlocal ssh
+        if ssh is not None:
+            return ssh
+        ssh = paramiko.SSHClient()
+        ssh.load_system_host_keys()
+        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        ssh.connect(args.host)
+        return ssh
+
+    sftp = None
+    def get_sftp():
+        nonlocal sftp
+        if sftp is not None:
+            return sftp
+        sftp = get_ssh().open_sftp()
+        return sftp
+
+    rng = Random()
+    tmp = "artiq" + "".join([rng.choice("ABCDEFGHIJKLMNOPQRSTUVWXYZ") for _ in range(6)])
+    env = "bash -c 'export PATH=$HOME/miniconda/bin:$PATH; exec $0 $*' "
+
+    def run_command(cmd):
+        logger.info("Executing {}".format(cmd))
+        chan = get_ssh().get_transport().open_session()
+        chan.set_combine_stderr(True)
+        chan.exec_command(cmd.format(tmp=tmp, env=env, serial=args.serial, ip=args.ip))
+        return chan.makefile()
+
+    def drain(chan):
+        while True:
+            char = chan.read(1)
+            if char == b"":
+                break
+            sys.stderr.write(char.decode("utf-8"))
+
+    for action in args.actions:
+        if action == "build":
+            logger.info("Building runtime")
+            subprocess.call(["python3", "-m", "artiq.gateware.targets.kc705",
+                                        "-H", "nist_clock",
+                                        "--no-compile-gateware",
+                                        "--output-dir", "/tmp/kc705"])
+
+        elif action == "boot":
+            logger.info("Uploading runtime")
+            get_sftp().mkdir("/tmp/{tmp}".format(tmp=tmp))
+            get_sftp().put("/tmp/kc705/software/runtime/runtime.bin",
+                           "/tmp/{tmp}/runtime.bin".format(tmp=tmp))
+
+            logger.info("Booting runtime")
+            flterm = run_command(
+                "{env} python3 flterm.py {serial} " +
+                "--kernel /tmp/{tmp}/runtime.bin --upload-only")
+            artiq_flash = run_command(
+                "{env} artiq_flash start")
+            drain(flterm)
+
+        elif action == "connect":
+            def forwarder(port):
+                listener = socket.socket()
+                listener.bind(('localhost', port))
+                listener.listen(1)
+                while True:
+                    local_stream, peer_addr = listener.accept()
+                    logger.info("Accepting %s:%s and opening SSH channel to %s:%s",
+                                *peer_addr, args.ip, port)
+                    remote_stream = get_ssh().get_transport() \
+                        .open_channel('direct-tcpip', (args.ip, port), peer_addr)
+                    while True:
+                        r, w, x = select.select([local_stream, remote_stream], [], [])
+                        if local_stream in r:
+                            data = local_stream.recv(1024)
+                            if data == b"":
+                                break
+                            remote_stream.send(data)
+                        if remote_stream in r:
+                            data = remote_stream.recv(1024)
+                            if data == b"":
+                                break
+                            local_stream.send(data)
+
+            for port in (1381, 1382):
+                thread = threading.Thread(target=forwarder, args=(port,),
+                                          name="port-{}".format(port), daemon=True)
+                thread.start()
+
+            logger.info("Connecting to device")
+            flterm = run_command(
+                "{env} python3 flterm.py {serial} --output-only")
+            drain(flterm)
+
+        else:
+            logger.error("Unknown action {}".format(action))
+            sys.exit(1)
+
+if __name__ == "__main__":
+    main()
diff --git a/setup.py b/setup.py
index c3c50e8c1..693da9a21 100755
--- a/setup.py
+++ b/setup.py
@@ -18,6 +18,7 @@ requirements = [
     "quamash", "pyqtgraph", "pygit2", "aiohttp",
     "llvmlite_artiq", "pythonparser", "python-Levenshtein",
     "lit", "OutputCheck",
+    "paramiko"
 ]
 
 console_scripts = [
@@ -27,6 +28,7 @@ console_scripts = [
     "artiq_coreconfig=artiq.frontend.artiq_coreconfig:main",
     "artiq_corelog=artiq.frontend.artiq_corelog:main",
     "artiq_ctlmgr=artiq.frontend.artiq_ctlmgr:main",
+    "artiq_devtool=artiq.frontend.artiq_devtool:main",
     "artiq_influxdb=artiq.frontend.artiq_influxdb:main",
     "artiq_master=artiq.frontend.artiq_master:main",
     "artiq_mkfs=artiq.frontend.artiq_mkfs:main",