forked from M-Labs/artiq
1
0
Fork 0

artiq_pcap: implement.

This commit is contained in:
whitequark 2017-06-25 07:04:29 +00:00
parent 69fa9b38e0
commit 816ec6c52f
5 changed files with 136 additions and 58 deletions

1
.gitignore vendored
View File

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

View File

@ -10,12 +10,10 @@ import subprocess
import socket import socket
import select import select
import threading import threading
import paramiko
import os import os
import shutil 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(): def get_argparser():
@ -60,43 +58,13 @@ def main():
else: else:
raise NotImplementedError("unknown target {}".format(args.target)) 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: for action in args.actions:
if action == "build": if action == "build":
@ -119,27 +87,29 @@ def main():
elif action == "reset": elif action == "reset":
logger.info("Resetting device") logger.info("Resetting device")
artiq_flash = run_command( client.run_command(
"{env} artiq_flash start" + "{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": elif action == "boot" or action == "boot+log":
logger.info("Uploading firmware") 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") logger.info("Booting firmware")
flterm = run_command( flterm = client.spawn_command(
"{env} python3 flterm.py {serial} " + "{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" + "{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": elif action == "connect":
def forwarder(port): def forwarder(port):
@ -151,12 +121,12 @@ def main():
local_stream, peer_addr = listener.accept() local_stream, peer_addr = listener.accept()
logger.info("Accepting %s:%s and opening SSH channel to %s:%s", logger.info("Accepting %s:%s and opening SSH channel to %s:%s",
*peer_addr, args.ip, port) *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!") logger.error("Trying to open a channel before the transport is ready!")
continue continue
try: try:
remote_stream = get_ssh().get_transport() \ remote_stream = client.get_transport() \
.open_channel('direct-tcpip', (args.ip, port), peer_addr) .open_channel('direct-tcpip', (args.ip, port), peer_addr)
except Exception as e: except Exception as e:
logger.exception("Cannot open channel on port %s", port) logger.exception("Cannot open channel on port %s", port)
@ -186,9 +156,9 @@ def main():
thread.start() thread.start()
logger.info("Connecting to device") 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": elif action == "hotswap":
logger.info("Hotswapping firmware") logger.info("Hotswapping firmware")

View File

@ -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))

View File

@ -6,8 +6,10 @@ import collections
import os import os
import atexit import atexit
import string import string
import random
import numpy as np import numpy as np
import paramiko
from artiq.language.environment import is_experiment from artiq.language.environment import is_experiment
from artiq.protocols import pyon from artiq.protocols import pyon
@ -253,3 +255,51 @@ def get_user_config_dir():
dir = user_config_dir("artiq", "m-labs", major) dir = user_config_dir("artiq", "m-labs", major)
os.makedirs(dir, exist_ok=True) os.makedirs(dir, exist_ok=True)
return dir 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))

View File

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