From d75054fb153078d7cab278c4e21736bc9d7c67f1 Mon Sep 17 00:00:00 2001
From: Vishnu Ks <yo@vishnuks.com>
Date: Fri, 13 Jul 2018 21:04:39 +0530
Subject: [PATCH] billing: Add function to sign strings.

---
 zilencer/lib/stripe.py        | 13 ++++++++++++-
 zilencer/tests/test_stripe.py | 12 +++++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/zilencer/lib/stripe.py b/zilencer/lib/stripe.py
index eb189999e7..45a51e33a8 100644
--- a/zilencer/lib/stripe.py
+++ b/zilencer/lib/stripe.py
@@ -2,17 +2,19 @@ import datetime
 from functools import wraps
 import logging
 import os
-from typing import Any, Callable, Optional, TypeVar
+from typing import Any, Callable, Optional, TypeVar, Tuple
 import ujson
 
 from django.conf import settings
 from django.db import transaction
 from django.utils.translation import ugettext as _
+from django.core.signing import Signer
 import stripe
 
 from zerver.lib.exceptions import JsonableError
 from zerver.lib.logging_util import log_to_file
 from zerver.lib.timestamp import datetime_to_timestamp, timestamp_to_datetime
+from zerver.lib.utils import generate_random_token
 from zerver.models import Realm, UserProfile, RealmAuditLog
 from zilencer.models import Customer, Plan
 from zproject.settings import get_secret
@@ -52,6 +54,15 @@ CallableT = TypeVar('CallableT', bound=Callable[..., Any])
 def get_seat_count(realm: Realm) -> int:
     return UserProfile.objects.filter(realm=realm, is_active=True, is_bot=False).count()
 
+def sign_string(string: str) -> Tuple[str, str]:
+    salt = generate_random_token(64)
+    signer = Signer(salt=salt)
+    return signer.sign(string), salt
+
+def unsign_string(signed_string: str, salt: str) -> str:
+    signer = Signer(salt=salt)
+    return signer.unsign(signed_string)
+
 class StripeError(JsonableError):
     pass
 
diff --git a/zilencer/tests/test_stripe.py b/zilencer/tests/test_stripe.py
index 22712bb08f..3e3ad3ea25 100644
--- a/zilencer/tests/test_stripe.py
+++ b/zilencer/tests/test_stripe.py
@@ -3,6 +3,8 @@ import os
 from typing import Any
 import ujson
 
+from django.core import signing
+
 import stripe
 from stripe.api_resources.list_object import ListObject
 
@@ -13,7 +15,7 @@ from zerver.lib.timestamp import timestamp_to_datetime
 from zerver.models import Realm, UserProfile, get_realm, RealmAuditLog
 from zilencer.lib.stripe import StripeError, catch_stripe_errors, \
     do_create_customer_with_payment_source, do_subscribe_customer_to_plan, \
-    get_seat_count, extract_current_subscription
+    get_seat_count, extract_current_subscription, sign_string, unsign_string
 from zilencer.models import Customer, Plan
 
 fixture_data_file = open(os.path.join(os.path.dirname(__file__), 'stripe_fixtures.json'), 'r')
@@ -246,6 +248,14 @@ class StripeTest(ZulipTestCase):
         subscription = extract_current_subscription(customer_with_subscription)
         self.assertEqual(subscription["id"][:4], "sub_")
 
+    def test_sign_string(self) -> None:
+        string = "abc"
+        signed_string, salt = sign_string(string)
+        self.assertEqual(string, unsign_string(signed_string, salt))
+
+        with self.assertRaises(signing.BadSignature):
+            unsign_string(signed_string, "randomsalt")
+
 class BillingUpdateTest(ZulipTestCase):
     def setUp(self) -> None:
         self.user = self.example_user("hamlet")