Implemet cookie encrypting

kobin/environs.py

@@ -22,7 +22,11 @@
 In addition to the :class:`Response` class, Kobin provides :class:`TemplateResponse` ,
  :class:`JSONResponse` , :class:`RedirectResponse` and :class:`HTTPError`.
 """
+import base64
+import hashlib
+import hmac
 import threading
+import time
 import cgi
 import json
 import http.client as http_client
@@ -34,6 +38,9 @@
 ##################################################################################
 # Request Object #################################################################
 ##################################################################################
+import pickle
+
+
 class Request:
     """ A wrapper for WSGI environment dictionaries.
     """
@@ -111,8 +118,18 @@ def cookies(self):
         cookies = SimpleCookie(self.environ.get('HTTP_COOKIE', '')).values()
         return {c.key: c.value for c in cookies}
 
-    def get_cookie(self, key, default=None, secret=None):
+    def get_cookie(self, key, default=None, secret=None, digestmod=hashlib.sha256):
         value = self.cookies.get(key)
+        if secret and value and value.startswith('!') and '?' in value:
+            # See BaseResponse.set_cookie for details.
+            if isinstance(secret, str):
+                secret = secret.encode('utf-8')
+            sig, msg = map(lambda x: x.encode('utf-8'), value[1:].split('?', 1))
+            hash_string = hmac.new(secret, msg, digestmod=digestmod).digest()
+            if sig == base64.b64encode(hash_string):
+                key_and_value = pickle.loads(base64.b64decode(msg))
+                if key_and_value and key_and_value[0] == key:
+                    return key_and_value[1]
         return value or default
 
     def __getitem__(self, key):
@@ -228,9 +245,20 @@ def headerlist(self):
                 self.headers.add_header('Set-Cookie', c.OutputString())
         return self.headers.items()
 
-    def set_cookie(self, key, value, expires=None, max_age=None, path=None):
-        import time
+    def set_cookie(self, key, value, expires=None, max_age=None, path=None,
+                   secret=None, digestmod=hashlib.sha256):
+        if secret:
+            if isinstance(secret, str):
+                secret = secret.encode('utf-8')
+            encoded = base64.b64encode(pickle.dumps((key, value), pickle.HIGHEST_PROTOCOL))
+            sig = base64.b64encode(hmac.new(secret, encoded, digestmod=digestmod).digest())
+            value_bytes = b'!' + sig + b'?' + encoded
+            value = value_bytes.decode('utf-8')
+
         self._cookies[key] = value
+        if len(key) + len(value) > 3800:
+            raise ValueError('Content does not fit into a cookie.')
+
         if max_age is not None:
             if isinstance(max_age, int):
                 max_age_value = max_age

kobin/environs.pyi

@@ -35,7 +35,7 @@ class Request:
     def url(self) -> str: ...
     @property
     def cookies(self) -> Dict[str, str]: ...
-    def get_cookie(self, key: str, default: str = ..., secret: str = ...) -> str: ...
+    def get_cookie(self, key: str, default: str = ..., secret: Union[str, bytes] = ...) -> str: ...
     def __getitem__(self, key: str): ...
     def __delitem__(self, key: str): ...
     def __setitem__(self, key: str, value: Any): ...
@@ -77,8 +77,10 @@ class BaseResponse:
     def status(self, status_code: int) -> None: ...
     @property
     def headerlist(self) -> List[Tuple[str, str]]: ...
-    def set_cookie(self, key: str, value: Any, expires: Union[date, datetime, int] = ...,
-                   max_age: Union[timedelta, int] = ..., path: str = ...) -> None: ...
+    def set_cookie(self, key: str, value: str,
+                   expires: Union[date, datetime, int] = ..., max_age: Union[timedelta, int] = ...,
+                   path: str = ..., secret: Union[str, bytes] = ...,
+                   digestmod: Callable[..., bytes] = ...) -> None: ...
     def delete_cookie(self, key: str, **kwargs: Any) -> None: ...
 
 class Response:

tests/test_environs.py

@@ -178,6 +178,7 @@ def test_get_cookie(self):
         expected = 'bar'
         self.assertEqual(actual, expected)
 
+    # Delete Cookie Tests in Request Class
     @freezegun.freeze_time('2017-01-01 00:00:00')
     def test_delete_cookie(self):
         response = BaseResponse()
@@ -187,6 +188,21 @@ def test_delete_cookie(self):
             'foo=""; expires=Sun, 01 Jan 2017 00:00:00 GMT; Max-Age=-1')
         self.assertIn(expected_set_cookie, response.headerlist)
 
+    # Get and Set Cookie Tests with secret
+    def test_set_cookie_with_secret(self):
+        response = BaseResponse()
+        response.set_cookie('foo', 'bar', secret='secretkey')
+        expected_set_cookie = ('Set-Cookie', 'foo="!VzhGFLGcW+5OMs1s4beLXaqFxAUwgHdWkH5fgapghoI='
+                                             '?gASVDwAAAAAAAACMA2Zvb5SMA2JhcpSGlC4="')
+        self.assertIn(expected_set_cookie, response.headerlist)
+
+    def test_get_cookie_with_secret(self):
+        request = Request({'HTTP_COOKIE': 'foo="!VzhGFLGcW+5OMs1s4beLXaqFxAUwgHdWkH5fgapghoI='
+                                          '?gASVDwAAAAAAAACMA2Zvb5SMA2JhcpSGlC4="'})
+        actual = request.get_cookie("foo", secret='secretkey')
+        expected = 'bar'
+        self.assertEqual(actual, expected)
+
 
 class BaseResponseTests(TestCase):
     def test_constructor_body(self):