diff --git a/.gitignore b/.gitignore index f04678d44..91ffc1c07 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,7 @@ __pycache__/ *.bin *.elf *.fbi +*.pcap .ipynb_checkpoints /doc/manual/_build /build diff --git a/artiq/frontend/artiq_devtool.py b/artiq/frontend/artiq_devtool.py index 90b6556e4..a2da3220d 100755 --- a/artiq/frontend/artiq_devtool.py +++ b/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 random import Random +from artiq.tools import verbosity_args, init_logger, logger, SSHClient def get_argparser(): @@ -60,43 +58,13 @@ def main(): else: raise NotImplementedError("unknown target {}".format(args.target)) - 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, **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')) + client = SSHClient(args.host) + substs = { + "env": "bash -c 'export PATH=$HOME/miniconda/bin:$PATH; exec $0 $*' ", + "serial": args.serial, + "ip": args.ip, + "firmware": firmware, + } 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 "")) - drain(artiq_flash) + (" --target-file " + args.config if args.config else ""), + **substs) elif action == "boot" or action == "boot+log": logger.info("Uploading firmware") - get_sftp().mkdir("/tmp/{tmp}".format(tmp=tmp)) - get_sftp().put("/tmp/{target}/software/{firmware}/{firmware}.bin" - .format(target=args.target, firmware=firmware), - "/tmp/{tmp}/{firmware}.bin".format(tmp=tmp, firmware=firmware)) + client.get_sftp().put("/tmp/{target}/software/{firmware}/{firmware}.bin" + .format(target=args.target, firmware=firmware), + "{tmp}/{firmware}.bin" + .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 " + - ("--upload-only" if action == "boot" else "--output-only")) - artiq_flash = run_command( + "--kernel {tmp}/{firmware}.bin " + + ("--upload-only" if action == "boot" else "--output-only"), + **substs) + artiq_flash = client.spawn_command( "{env} artiq_flash start" + - (" --target-file " + args.config if args.config else "")) - drain(flterm) + (" --target-file " + args.config if args.config else ""), + **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( - "{env} python3 flterm.py {serial} --output-only") - drain(flterm) + client.run_command( + "{env} python3 flterm.py {serial} --output-only", + **substs) elif action == "hotswap": logger.info("Hotswapping firmware") diff --git a/artiq/frontend/artiq_pcap.py b/artiq/frontend/artiq_pcap.py new file mode 100644 index 000000000..26db1a658 --- /dev/null +++ b/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)) diff --git a/artiq/tools.py b/artiq/tools.py index bd90e7902..352647d6b 100644 --- a/artiq/tools.py +++ b/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)) diff --git a/setup.py b/setup.py index 8f36657bd..243ea4fe7 100755 --- a/setup.py +++ b/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",