artiq_pcap: implement.

.gitignore

@@ -6,6 +6,7 @@ __pycache__/
 *.bin
 *.elf
 *.fbi
+*.pcap
 .ipynb_checkpoints
 /doc/manual/_build
 /build

artiq/frontend/artiq_devtool.py

@@ -10,12 +10,10 @@ import subprocess
 import socket
 import select
 import threading
-import paramiko
 import os
 import shutil
 
-from artiq.tools import verbosity_args, init_logger, logger
+from artiq.tools import verbosity_args, init_logger, logger, SSHClient
-from random import Random
 
 
 def get_argparser():
@@ -60,43 +58,13 @@ def main():
     else:
         raise NotImplementedError("unknown target {}".format(args.target))
 
-    ssh = None
+    client = SSHClient(args.host)
-    def get_ssh():
+    substs = {
-        nonlocal ssh
+        "env":      "bash -c 'export PATH=$HOME/miniconda/bin:$PATH; exec $0 $*' ",
-        if ssh is not None:
+        "serial":   args.serial,
-            return ssh
+        "ip":       args.ip,
-        ssh = paramiko.SSHClient()
+        "firmware": firmware,
-        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, **kws):
-        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,
-                                     firmware=firmware, **kws))
-        return chan.makefile()
-
-    def drain(chan):
-        while True:
-            char = chan.read(1)
-            if char == b"":
-                break
-            sys.stderr.write(char.decode("utf-8", errors='replace'))
 
     for action in args.actions:
         if action == "build":
@@ -119,27 +87,29 @@ def main():
 
         elif action == "reset":
             logger.info("Resetting device")
-            artiq_flash = run_command(
+            client.run_command(
                 "{env} artiq_flash start" +
-                (" --target-file " + args.config if args.config else ""))
+                (" --target-file " + args.config if args.config else ""),
-            drain(artiq_flash)
+                **substs)
 
         elif action == "boot" or action == "boot+log":
             logger.info("Uploading firmware")
-            get_sftp().mkdir("/tmp/{tmp}".format(tmp=tmp))
+            client.get_sftp().put("/tmp/{target}/software/{firmware}/{firmware}.bin"
-            get_sftp().put("/tmp/{target}/software/{firmware}/{firmware}.bin"
+                                      .format(target=args.target, firmware=firmware),
-                                .format(target=args.target, firmware=firmware),
+                                  "{tmp}/{firmware}.bin"
-                           "/tmp/{tmp}/{firmware}.bin".format(tmp=tmp, firmware=firmware))
+                                      .format(tmp=client.tmp, firmware=firmware))
 
             logger.info("Booting firmware")
-            flterm = run_command(
+            flterm = client.spawn_command(
                 "{env} python3 flterm.py {serial} " +
-                "--kernel /tmp/{tmp}/{firmware}.bin " +
+                "--kernel {tmp}/{firmware}.bin " +
-                ("--upload-only" if action == "boot" else "--output-only"))
+                ("--upload-only" if action == "boot" else "--output-only"),
-            artiq_flash = run_command(
+                **substs)
+            artiq_flash = client.spawn_command(
                 "{env} artiq_flash start" +
-                (" --target-file " + args.config if args.config else ""))
+                (" --target-file " + args.config if args.config else ""),
-            drain(flterm)
+                **substs)
+            client.drain(flterm)
 
         elif action == "connect":
             def forwarder(port):
@@ -151,12 +121,12 @@ def main():
                     local_stream, peer_addr = listener.accept()
                     logger.info("Accepting %s:%s and opening SSH channel to %s:%s",
                                 *peer_addr, args.ip, port)
-                    if get_ssh().get_transport() is None:
+                    if client.get_transport() is None:
                         logger.error("Trying to open a channel before the transport is ready!")
                         continue
 
                     try:
-                        remote_stream = get_ssh().get_transport() \
+                        remote_stream = client.get_transport() \
                             .open_channel('direct-tcpip', (args.ip, port), peer_addr)
                     except Exception as e:
                         logger.exception("Cannot open channel on port %s", port)
@@ -186,9 +156,9 @@ def main():
                 thread.start()
 
             logger.info("Connecting to device")
-            flterm = run_command(
+            client.run_command(
-                "{env} python3 flterm.py {serial} --output-only")
+                "{env} python3 flterm.py {serial} --output-only",
-            drain(flterm)
+                **substs)
 
         elif action == "hotswap":
             logger.info("Hotswapping firmware")

artiq/frontend/artiq_pcap.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+
+# This script makes the following assumptions:
+#  * tcpdump has CAP_NET_RAW capabilities set
+#    use # setcap cap_net_raw+eip /usr/sbin/tcpdump
+
+import argparse
+import subprocess
+
+from artiq.tools import verbosity_args, init_logger, logger, SSHClient
+
+
+def get_argparser():
+    parser = argparse.ArgumentParser(description="ARTIQ core device "
+                                                 "packet capture tool")
+
+    verbosity_args(parser)
+
+    parser.add_argument("--host", metavar="HOST",
+                        type=str, default="lab.m-labs.hk",
+                        help="SSH host where the development board is located")
+    parser.add_argument("-i", "--ip", metavar="IP",
+                        type=str, default="kc705.lab.m-labs.hk",
+                        help="IP address corresponding to the development board")
+    parser.add_argument("-f", "--file", metavar="PCAP_FILE",
+                        type=str, default="coredevice.pcap",
+                        help="Location to retrieve the pcap file into")
+
+    parser.add_argument("command", metavar="COMMAND",
+                        type=str, default=[], nargs="+",
+                        help="command to execute while capturing")
+
+    return parser
+
+
+def main():
+    args = get_argparser().parse_args()
+    init_logger(args)
+
+    client = SSHClient(args.host)
+
+    sftp = client.get_sftp()
+    tcpdump = client.spawn_command(
+        "/usr/sbin/tcpdump host {ip} -w {tmp}/trace.pcap", get_pty=True,
+        ip=args.ip)
+
+    try:
+        subprocess.check_call(args.command)
+    except subprocess.CalledProcessError:
+        logger.error("Command failed")
+        sys.exit(1)
+
+    tcpdump.close()
+    sftp.get("{tmp}/trace.pcap".format(tmp=client.tmp),
+             args.file)
+    logger.info("Pcap file {file} retrieved".format(file=args.file))

artiq/tools.py

@@ -6,8 +6,10 @@ import collections
 import os
 import atexit
 import string
+import random
 
 import numpy as np
+import paramiko
 
 from artiq.language.environment import is_experiment
 from artiq.protocols import pyon
@@ -253,3 +255,51 @@ def get_user_config_dir():
     dir = user_config_dir("artiq", "m-labs", major)
     os.makedirs(dir, exist_ok=True)
     return dir
+
+
+class SSHClient:
+    def __init__(self, host):
+        self.host = host
+        self.ssh = None
+        self.sftp = None
+
+        tmpname = "".join([random.Random().choice("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
+                           for _ in range(6)])
+        self.tmp = "/tmp/artiq" + tmpname
+
+    def get_ssh(self):
+        if self.ssh is None:
+            self.ssh = paramiko.SSHClient()
+            self.ssh.load_system_host_keys()
+            self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+            self.ssh.connect(self.host)
+        return self.ssh
+
+    def get_transport(self):
+        return self.get_ssh().get_transport()
+
+    def get_sftp(self):
+        if self.sftp is None:
+            self.sftp = self.get_ssh().open_sftp()
+            self.sftp.mkdir(self.tmp)
+            atexit.register(lambda: self.run_command("rm -rf {tmp}"))
+        return self.sftp
+
+    def spawn_command(self, cmd, get_pty=False, **kws):
+        logger.info("Executing {}".format(cmd))
+        chan = self.get_ssh().get_transport().open_session()
+        if get_pty:
+            chan.get_pty()
+        chan.set_combine_stderr(True)
+        chan.exec_command(cmd.format(tmp=self.tmp, **kws))
+        return chan
+
+    def drain(self, chan):
+        while True:
+            char = chan.recv(1)
+            if char == b"":
+                break
+            sys.stderr.write(char.decode("utf-8", errors='replace'))
+
+    def run_command(self, cmd, **kws):
+        self.drain(self.spawn_command(cmd, **kws))

setup.py

@@ -27,6 +27,7 @@ console_scripts = [
     "artiq_coreboot = artiq.frontend.artiq_coreboot:main",
     "artiq_ctlmgr = artiq.frontend.artiq_ctlmgr:main",
     "artiq_devtool = artiq.frontend.artiq_devtool:main",
+    "artiq_pcap = artiq.frontend.artiq_pcap:main",
     "artiq_influxdb = artiq.frontend.artiq_influxdb:main",
     "artiq_master = artiq.frontend.artiq_master:main",
     "artiq_mkfs = artiq.frontend.artiq_mkfs:main",