add AFWS client

Author: sbourdeauducq

MANIFEST.in

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

artiq/afws.pem

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

artiq/frontend/afws_client.py

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

setup.py

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