feat: add support for custom kms hosts (#5)

* feat: add support for custom kms hosts

* chore: update tests

* chore: bump version

Author: rohan-chaturvedi

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "phase_dev"
-version = "1.0.0"
+version = "1.1.0"
 description = "Python SDK for Phase"
 readme = "README.md"
 requires-python = ">=3.10"

src/phase/phase.py

@@ -4,6 +4,8 @@
 from .utils.crypto import decrypt_b64, encrypt_b64, fetch_app_key, random_key_pair, reconstruct_secret
 from .version import __version__, __ph_version__
 
+DEFAULT_KMS_HOST = "https://kms.phase.dev"
+
 
 @dataclass
 class AppSecret:
@@ -18,8 +20,9 @@ class Phase:
     _app_id = ''
     _app_pub_key = ''
     _app_secret = None
+    _kms_host = ''
 
-    def __init__(self, app_id, app_secret):
+    def __init__(self, app_id, app_secret, custom_kms_host=None):
         app_id_pattern = re.compile(r"^phApp:v(\d+):([a-fA-F0-9]{64})$")
         app_secret_pattern = re.compile(
             r"^pss:v(\d+):([a-fA-F0-9]{64}):([a-fA-F0-9]{64,128}):([a-fA-F0-9]{64})$")
@@ -32,9 +35,10 @@ def __init__(self, app_id, app_secret):
 
         self._app_id = app_id
         self._app_pub_key = app_id.split(':')[2]
-        app_secret_segments = app_secret.split(':')
 
+        app_secret_segments = app_secret.split(':')
         self._app_secret = AppSecret(*app_secret_segments)
+        self._kms_host = f"{custom_kms_host}/kms" if custom_kms_host else DEFAULT_KMS_HOST
 
     def encrypt(self, plaintext, tag="") -> str | None:
         """
@@ -81,7 +85,7 @@ def decrypt(self, phase_ciphertext) -> str | None:
             client_pub_key = bytes.fromhex(client_pub_key_hex)
 
             keyshare1 = fetch_app_key(
-                self._app_secret.app_token, self._app_secret.keyshare1_unwrap_key, self._app_id, len(ct)/2)
+                self._app_secret.app_token, self._app_secret.keyshare1_unwrap_key, self._app_id, len(ct)/2, self._kms_host)
 
             app_priv_key = reconstruct_secret(
                 [self._app_secret.keyshare0, keyshare1])

src/phase/utils/crypto.py

@@ -6,9 +6,6 @@
 from ..version import __version__
 
 
-PHASE_KMS_URI = "https://kms.phase.dev/"
-
-
 def xor_bytes(a, b) -> bytes:
     """
     Computes the XOR of two byte arrays byte by byte.
@@ -129,7 +126,7 @@ def decrypt_b64(ct, key) -> bytes:
     return plaintext_bytes.decode('utf-8')
 
 
-def fetch_app_key(appToken, wrapKey, appId, dataSize) -> str:
+def fetch_app_key(appToken, wrapKey, appId, dataSize, host) -> str:
     """
     Fetches the application key share from Phase KMS.
 
@@ -152,7 +149,7 @@ def fetch_app_key(appToken, wrapKey, appId, dataSize) -> str:
         "PhSize": f"{dataSize}"
     }
 
-    response = requests.get(f"{PHASE_KMS_URI}{appId}", headers=headers)
+    response = requests.get(f"{host}/{appId}", headers=headers)
 
     if response.status_code == 404:
         raise Exception("Invalid app token")

src/phase/version.py

@@ -1,2 +1,2 @@
-__version__ = "1.0.0"
+__version__ = "1.1.0"
 __ph_version__ = "v1"

tests/test_decrypt.py

@@ -11,7 +11,7 @@ def phase_instance():
     return Phase(APP_ID, APP_SECRET)
 
 
-def mock_fetch_app_key(appToken, wrapKey, appId, dataSize):
+def mock_fetch_app_key(appToken, wrapKey, appId, dataSize, custom_kms_host=None):
     return "e35ae9560207c90fa3dd68a8715e13a1ef988bffa284db73f04328df17f37cfe"