diff --git a/bin/require.pip b/bin/require.pip index d691a00e5d..e9f2315b2d 100644 --- a/bin/require.pip +++ b/bin/require.pip @@ -3,3 +3,4 @@ GitPython==2.1.3 requests==2.18.2 sarge==0.1.4 pytest==3.2.1 +psycopg2==2.7.5 diff --git a/test/server/clientlib.py b/test/server/clientlib.py index b02888b0ea..2541893d0c 100644 --- a/test/server/clientlib.py +++ b/test/server/clientlib.py @@ -7,6 +7,7 @@ import uuid import random import time +from pglib import attach_device example_images = {} execfile(os.path.normpath(os.path.join(__file__, "../../../bin/load_test_exercise_images.py")), example_images) @@ -20,6 +21,7 @@ def __init__(self, backend="http://localhost:10080"): self.deviceInfo = make_device_info() self.deviceId = make_uuid() self.secret = make_uuid() + self.accountId = make_random_id() self.session = requests.Session() self.session.headers.update({'Accept-Language': 'en-US'}) @@ -31,6 +33,8 @@ def login(self): resp = self.session.post( urljoin(self.backend, "/api/register"), data=dict(deviceId=self.deviceId, secret=self.secret, deviceInfo=json.dumps(self.deviceInfo))) + account_info = attach_device(self.deviceId, self.accountId) + self.session.cookies.update(account_info) resp.raise_for_status() return resp diff --git a/test/server/pglib.py b/test/server/pglib.py new file mode 100644 index 0000000000..90475eedcf --- /dev/null +++ b/test/server/pglib.py @@ -0,0 +1,42 @@ +# This is for faking an FxA auth by directly updating the database. FxA is +# required for setting a shot's expiration, which is in multiple tests. + +import os +import subprocess +import shlex +import psycopg2 +import uuid +import hmac +import hashlib +from base64 import urlsafe_b64encode + +dir_path = os.path.dirname(os.path.realpath(__file__)) +pg_vars_path = os.path.join(dir_path, "..", "..", "bin", "pg_vars") +pg_vars_out = subprocess.check_output(pg_vars_path) +parts = [line.partition('=') for line in shlex.split(pg_vars_out)] +pg_vars = {name: val for name, _, val in parts} + + +# Attaches a device_id to an account_id. Returns a dict that can be stuffed +# into a cookie jar. +def attach_device(device_id, account_id): + dbname = pg_vars["PGDATABASE"] or pg_vars["PGUSER"] + conn = psycopg2.connect(dbname=dbname, user=pg_vars["PGUSER"], host=pg_vars["PGHOST"], port=pg_vars["PGPORT"]) + conn.autocommit = True + cur = conn.cursor() + token = str(uuid.uuid1()) + cur.execute("INSERT INTO accounts (id, token) VALUES (%s, %s)", (account_id, token)) + cur.execute("UPDATE devices SET accountid = %s WHERE id = %s", (account_id, device_id)) + cur.execute("SELECT key FROM signing_keys ORDER BY created DESC LIMIT 1") + key_row = cur.fetchone() + account_id_hmac = __get_hmac("accountid=%s" % account_id, key_row[0]) + cur.close() + conn.close() + return {"accountid": account_id, "accountid.sig": account_id_hmac} + + +def __get_hmac(val, key): + h = hmac.new(key, None, hashlib.sha1) + h.update(val) + b64 = urlsafe_b64encode(h.digest()) + return b64.replace('=', '')