add AFWS client

This commit is contained in:
Sebastien Bourdeauducq 2022-02-07 14:28:00 +08:00
parent d5eec652ee
commit 833acb6925
4 changed files with 191 additions and 0 deletions

View File

@ -5,3 +5,4 @@ include versioneer.py
include artiq/_version.py include artiq/_version.py
include artiq/coredevice/coredevice_generic.schema.json include artiq/coredevice/coredevice_generic.schema.json
include artiq/compiler/kernel.ld include artiq/compiler/kernel.ld
include artiq/afws.pem

23
artiq/afws.pem Normal file
View File

@ -0,0 +1,23 @@
-----BEGIN CERTIFICATE-----
MIID0zCCArugAwIBAgIUPkNfEUx/uau3z8SD4mgMbCK/DEgwDQYJKoZIhvcNAQEL
BQAweTELMAkGA1UEBhMCSEsxEzARBgNVBAgMClNvbWUtU3RhdGUxFzAVBgNVBAoM
Dk0tTGFicyBMaW1pdGVkMRkwFwYDVQQDDBBuaXhibGQubS1sYWJzLmhrMSEwHwYJ
KoZIhvcNAQkBFhJoZWxwZGVza0BtLWxhYnMuaGswHhcNMjIwMjA2MTA1ODQ0WhcN
MjUwMjA1MTA1ODQ0WjB5MQswCQYDVQQGEwJISzETMBEGA1UECAwKU29tZS1TdGF0
ZTEXMBUGA1UECgwOTS1MYWJzIExpbWl0ZWQxGTAXBgNVBAMMEG5peGJsZC5tLWxh
YnMuaGsxITAfBgkqhkiG9w0BCQEWEmhlbHBkZXNrQG0tbGFicy5oazCCASIwDQYJ
KoZIhvcNAQEBBQADggEPADCCAQoCggEBAPWetZhoggPR2ae7waGzv1AQ8NQO3noW
8DofVjusNpX5i/YB0waAr1bm1tALLJoHV2r/gTxujlXCe/L/WG1DLseCf6NO9sHg
t0FLhDpF9kPMWBgauVVLepd2Y2yU1G8eFuEVGnsiQSu0IzsZP5FQBJSyxvxJ+V/L
EW9ox91VGOP9VZR9jqdlYjGhcwClHA/nHe0q1fZq42+9rG466I5yIlNSoa7ilhTU
2C2doxy6Sr6VJYnLEMQqoIF65t3MkKi9iaqN7az0OCrj6XR0P5iKBzUhIgMUd2qs
7Id0XUdbQvaoaRI67vhGkNr+f4rdAUNCDGcbbokuBnmE7/gva6BAABUCAwEAAaNT
MFEwHQYDVR0OBBYEFM2e2FmcytXhKyfC1KEjVJ2mPSy3MB8GA1UdIwQYMBaAFM2e
2FmcytXhKyfC1KEjVJ2mPSy3MA8GA1UdEwEB/wQFMAMBAf8wDQYJKoZIhvcNAQEL
BQADggEBAKH0z5vlbfTghjYWwd2yEEFBbZx5XxaLHboFQpFpxu9sZoidVs047tco
MOr1py9juiNGGM8G35sw9306f+thDFwqlQfSExUwp5pRQNq+mxglMSF05HWDqBwb
wnItKi/WXpkMQXgpQJFVeflz4B4ZFNlH1UQl5bwacXOM9NM9zO7duCjVXmGE0yxi
VQyApfPQYu9whCSowDYYaA0toJeikMzGfWxhlAH79/2Qmit8KcSCbX1fK/QoRZLa
5NeUi/OlJbBpkgTrfzfMLphmsPWPAVMeUKzqd/vXfG6ZBOZZm6e6sl8RBycBezII
15WekikTE5+T54/E0xiu+zIW/Xhhk14=
-----END CERTIFICATE-----

166
artiq/frontend/afws_client.py Executable file
View File

@ -0,0 +1,166 @@
#!/usr/bin/env python3
import sys
import argparse
import os
import socket
import ssl
import io
import zipfile
from getpass import getpass
def get_artiq_cert():
try:
import artiq
except ImportError:
return None
filename = os.path.join(os.path.dirname(artiq.__file__), "afws.pem")
if not os.path.isfile(filename):
return None
return filename
def get_artiq_rev():
try:
import artiq
except ImportError:
return None
version = artiq.__version__
if version.endswith(".beta"):
version = version[:-5]
version = version.split(".")
if len(version) != 3:
return None
major, minor, rev = version
return rev
def zip_unarchive(data, directory):
buf = io.BytesIO(data)
with zipfile.ZipFile(buf) as archive:
archive.extractall(directory)
class Client:
def __init__(self, server, port, cafile):
self.ssl_context = ssl.create_default_context(cafile=cafile)
self.raw_socket = socket.create_connection((server, port))
try:
self.socket = self.ssl_context.wrap_socket(self.raw_socket, server_hostname=server)
except:
self.raw_socket.close()
raise
self.fsocket = self.socket.makefile("rwb")
def close(self):
self.socket.close()
self.raw_socket.close()
def send_command(self, *command):
self.fsocket.write((" ".join(command) + "\n").encode())
self.fsocket.flush()
def read_reply(self):
return self.fsocket.readline().decode("ascii").split()
def login(self, username, password):
self.send_command("LOGIN", username, password)
return self.read_reply() == ["HELLO"]
def build(self, rev, variant):
self.send_command("BUILD", rev, variant)
reply = self.read_reply()[0]
if reply != "BUILDING":
return reply, None
print("Build in progress. This may take 10-15 minutes.")
reply, status = self.read_reply()
if reply != "DONE":
raise ValueError("Unexpected server reply: expected 'DONE', got '{}'".format(reply))
if status != "done":
return status, None
print("Build completed. Downloading...")
reply, length = self.read_reply()
if reply != "PRODUCT":
raise ValueError("Unexpected server reply: expected 'PRODUCT', got '{}'".format(reply))
contents = self.fsocket.read(int(length))
print("Download completed.")
return "OK", contents
def passwd(self, password):
self.send_command("PASSWD", password)
return self.read_reply() == ["OK"]
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--server", default="nixbld.m-labs.hk", help="server to connect to (default: %(default)s)")
parser.add_argument("--port", default=7402, type=int, help="port to connect to (default: %(default)d)")
parser.add_argument("--cert", default=None, help="SSL certificate file used to authenticate server (default: afws.pem in ARTIQ)")
parser.add_argument("username", help="user name for logging into AFWS")
action = parser.add_subparsers(dest="action")
action.required = True
act_build = action.add_parser("build", help="build and download firmware")
act_build.add_argument("--rev", default=None, help="revision to build (default: currently installed ARTIQ revision)")
act_build.add_argument("variant", help="variant to build")
act_build.add_argument("directory", help="output directory")
act_passwd = action.add_parser("passwd", help="change password")
args = parser.parse_args()
cert = args.cert
if cert is None:
cert = get_artiq_cert()
if cert is None:
print("SSL certificate not found in ARTIQ. Specify manually using --cert.")
sys.exit(1)
if args.action == "passwd":
password = getpass("Current password: ")
else:
password = getpass()
client = Client(args.server, args.port, cert)
try:
if not client.login(args.username, password):
print("Login failed")
sys.exit(1)
if args.action == "passwd":
print("Password must made of alphanumeric characters (a-z, A-Z, 0-9) and be at least 8 characters long.")
password = getpass("New password: ")
password_confirm = getpass("New password (again): ")
while password != password_confirm:
print("Passwords do not match")
password = getpass("New password: ")
password_confirm = getpass("New password (again): ")
if not client.passwd(password):
print("Failed to change password")
sys.exit(1)
elif args.action == "build":
try:
os.mkdir(args.directory)
except FileExistsError:
if any(os.scandir(args.directory)):
print("Output directory already exists and is not empty. Please remove it and try again.")
sys.exit(1)
rev = args.rev
if rev is None:
rev = get_artiq_rev()
if rev is None:
print("Unable to determine currently installed ARTIQ revision. Specify manually using --rev.")
sys.exit(1)
result, contents = client.build(rev, args.variant)
if result != "OK":
if result == "UNAUTHORIZED":
print("You are not authorized to build this variant. Your firmware subscription may have expired. Contact helpdesk\x40m-labs.hk.")
else:
print("Build failed: {}".format(result))
sys.exit(1)
zip_unarchive(contents, args.directory)
else:
raise ValueError
finally:
client.close()
if __name__ == "__main__":
main()

View File

@ -33,6 +33,7 @@ console_scripts = [
"artiq_run = artiq.frontend.artiq_run:main", "artiq_run = artiq.frontend.artiq_run:main",
"artiq_flash = artiq.frontend.artiq_flash:main", "artiq_flash = artiq.frontend.artiq_flash:main",
"aqctl_corelog = artiq.frontend.aqctl_corelog:main", "aqctl_corelog = artiq.frontend.aqctl_corelog:main",
"afws_client = artiq.frontend.afws_client:main",
] ]
gui_scripts = [ gui_scripts = [