M-Labs/artiq
2
0
Fork 2
mirror of https://github.com/m-labs/artiq.git synced 2024-12-28 12:48:26 +08:00
Code Issues Projects Releases Wiki Activity

artiq_pcap: implement.

Browse Source
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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
.gitignore vendored
View File

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

86
artiq/frontend/artiq_devtool.py
View File

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

56
artiq/frontend/artiq_pcap.py Normal file
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))

50
artiq/tools.py
View File

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

1
setup.py
View File

@ -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",