From b4bddacedbaf9bc215990f7a2847abd85e02c19c Mon Sep 17 00:00:00 2001
From: Joseph Schorr <josephschorr@users.noreply.github.com>
Date: Tue, 12 Jan 2016 14:08:41 -0500
Subject: [PATCH] Switch to Fernet crypto as per gtank's recommendation

---
 config.py                 |  2 +-
 endpoints/api/__init__.py | 19 +++++++------------
 util/security/crypto.py   | 18 ++++++++++++++++++
 3 files changed, 26 insertions(+), 13 deletions(-)
 create mode 100644 util/security/crypto.py

diff --git a/config.py b/config.py
index 7dd862d80..d0766680e 100644
--- a/config.py
+++ b/config.py
@@ -293,4 +293,4 @@ class DefaultConfig(object):
   # "Secret" key for generating encrypted paging tokens. Only needed to be secret to
   # hide the ID range for production (in which this value is overridden). Should *not*
   # be relied upon for secure encryption otherwise.
-  PAGE_TOKEN_KEY = 'GdE4<~8&9LHDPM++'
+  PAGE_TOKEN_KEY = 'um=/?Kqgp)2yQaS/A6C{NL=dXE&>C:}('
diff --git a/endpoints/api/__init__.py b/endpoints/api/__init__.py
index 575e49a0c..2dc28490c 100644
--- a/endpoints/api/__init__.py
+++ b/endpoints/api/__init__.py
@@ -22,7 +22,7 @@ from auth.auth import process_oauth
 from endpoints.csrf import csrf_protect
 from endpoints.decorators import check_anon_protection
 from util.saas.metricqueue import time_decorator
-from util.security.aes import AESCipher
+from util.security.crypto import encrypt_string, decrypt_string
 
 
 logger = logging.getLogger(__name__)
@@ -220,26 +220,21 @@ def page_support(func):
   @query_param('next_page', 'The page token for the next page', type=str)
   def wrapper(self, query_args, *args, **kwargs):
     page_token = None
-    unecrypted = None
 
     if query_args['next_page']:
       # Decrypt the page token.
-      cipher = AESCipher(app.config['PAGE_TOKEN_KEY'])
-      try:
-        unecrypted = cipher.decrypt(query_args['next_page'])
-      except TypeError:
-        pass
-
-      if unecrypted is not None:
+      unencrypted = decrypt_string(query_args['next_page'], app.config['PAGE_TOKEN_KEY'])
+      if unencrypted is not None:
         try:
-          page_token = json.loads(unecrypted)
+          page_token = json.loads(unencrypted)
         except ValueError:
           pass
 
+    # Note: if page_token is None, we'll receive the first page of results back.
     (result, next_page_token) = func(self, query_args, page_token, *args, **kwargs)
     if next_page_token is not None:
-      cipher = AESCipher(app.config['PAGE_TOKEN_KEY'])
-      result['next_page'] = cipher.encrypt(json.dumps(next_page_token))
+      result['next_page'] = encrypt_string(json.dumps(next_page_token),
+                                           app.config['PAGE_TOKEN_KEY'])
 
     return result
 
diff --git a/util/security/crypto.py b/util/security/crypto.py
new file mode 100644
index 000000000..fe1de2953
--- /dev/null
+++ b/util/security/crypto.py
@@ -0,0 +1,18 @@
+import base64
+
+from cryptography.fernet import Fernet, InvalidToken
+
+def encrypt_string(string, key):
+  """ Encrypts a string with the specified key. The key must be 32 raw bytes. """
+  f = Fernet(base64.urlsafe_b64encode(key))
+  return f.encrypt(string)
+
+def decrypt_string(string, key):
+  """ Decrypts an encrypted string with the specified key. The key must be 32 raw bytes. """
+  f = Fernet(base64.urlsafe_b64encode(key))
+  try:
+    return f.decrypt(string)
+  except InvalidToken:
+    return None
+  except TypeError:
+    return None