srenblad
/
artiq
forked from M-Labs/artiq
1
0
Fork 0
Code Pull Requests Activity

afws_client: WebSocket, system certificates

This commit is contained in:
Sebastien Bourdeauducq 2023-04-07 16:03:33 +08:00
parent 97161a3df2
commit dc3db8bb66
1 changed files with 18 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
artiq/frontend/afws_client.py
View File

@ -13,17 +13,6 @@ from getpass import getpass
from tqdm import tqdm from tqdm import tqdm
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(): def get_artiq_rev():
try: try:
import artiq import artiq
@ -54,6 +43,7 @@ class Client:
def __init__(self, server, port, cafile): def __init__(self, server, port, cafile):
self.ssl_context = ssl.create_default_context(cafile=cafile) self.ssl_context = ssl.create_default_context(cafile=cafile)
self.raw_socket = socket.create_connection((server, port)) self.raw_socket = socket.create_connection((server, port))
self.init_websocket(server)
try: try:
self.socket = self.ssl_context.wrap_socket(self.raw_socket, server_hostname=server) self.socket = self.ssl_context.wrap_socket(self.raw_socket, server_hostname=server)
except: except:
@ -61,6 +51,19 @@ class Client:
raise raise
self.fsocket = self.socket.makefile("rwb") self.fsocket = self.socket.makefile("rwb")
def init_websocket(self, server):
self.raw_socket.sendall("GET / HTTP/1.1\r\nHost: {}\r\nConnection: Upgrade\r\nUpgrade: websocket\r\n\r\n"
.format(server).encode())
crlf_count = 0
while crlf_count < 4:
char = self.raw_socket.recv(1)
if not char:
return ValueError("Connection closed during WebSocket initialization")
if char == b"\r" or char == b"\n":
crlf_count += 1
else:
crlf_count = 0
def close(self): def close(self):
self.socket.close() self.socket.close()
self.raw_socket.close() self.raw_socket.close()
@ -160,9 +163,9 @@ class Client:
def main(): def main():
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("--server", default="nixbld.m-labs.hk", help="server to connect to (default: %(default)s)") parser.add_argument("--server", default="afws.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("--port", default=80, 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("--cert", default=None, help="SSL certificate file used to authenticate server (default: use system certificates)")
parser.add_argument("username", help="user name for logging into AFWS") parser.add_argument("username", help="user name for logging into AFWS")
action = parser.add_subparsers(dest="action") action = parser.add_subparsers(dest="action")
action.required = True action.required = True
@ -181,19 +184,12 @@ def main():
act_get_json.add_argument("-f", "--force", action="store_true", help="overwrite file if it already exists") act_get_json.add_argument("-f", "--force", action="store_true", help="overwrite file if it already exists")
args = parser.parse_args() 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": if args.action == "passwd":
password = getpass("Current password: ") password = getpass("Current password: ")
else: else:
password = getpass() password = getpass()
client = Client(args.server, args.port, cert) client = Client(args.server, args.port, args.cert)
try: try:
if not client.login(args.username, password): if not client.login(args.username, password):
print("Login failed") print("Login failed")