From 1f5c72ab14dccba00c41315bacb36ee21ecabdcb Mon Sep 17 00:00:00 2001
From: Barry Chen <chenba@gmail.com>
Date: Fri, 10 Aug 2018 16:51:11 -0500
Subject: [PATCH] Fix server tests that require an account id. (#4688)

---
 bin/require.pip          |  1 +
 test/server/clientlib.py |  4 ++++
 test/server/pglib.py     | 42 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+)
 create mode 100644 test/server/pglib.py

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('=', '')