refactor: restructure vault gRPC clients by version

Author: PromiseFru

Makefile

@@ -13,24 +13,28 @@ define download-proto
 	$(call log_message,INFO - $@ downloaded successfully!)
 endef
 
-$(PROTO_DIR)/%.proto:
+protos/v%/vault.proto:
 	$(eval PROTO_URL := $(PROTO_URL))
 	$(call download-proto)
 
 vault-proto: 
-	@rm -f "$(PROTO_DIR)/vault.proto"
-	@$(MAKE) PROTO_URL=https://raw.githubusercontent.com/smswithoutborders/RelaySMS-Vault/$(CURRENT_BRANCH)/protos/v1/vault.proto \
-		$(PROTO_DIR)/vault.proto
+	@for v in v1 v2; do \
+		rm -f "protos/$$v/vault.proto"; \
+		$(MAKE) PROTO_URL=https://raw.githubusercontent.com/smswithoutborders/RelaySMS-Vault/$(CURRENT_BRANCH)/protos/$$v/vault.proto \
+		protos/$$v/vault.proto; \
+	done
 
 grpc-compile:
-	$(call log_message,INFO - Compiling gRPC protos ...)
-	@$(python) -m grpc_tools.protoc \
-		-I$(PROTO_DIR) \
-		--python_out=. \
-		--pyi_out=. \
-		--grpc_python_out=. \
-		$(PROTO_DIR)/*.proto
-	$(call log_message,INFO - gRPC Compilation complete!)
+	$(call log_message,[INFO] Compiling gRPC protos ...)
+	@for v in v1 v2; do \
+		$(python) -m grpc_tools.protoc \
+			--proto_path=. \
+			--python_out=. \
+			--pyi_out=. \
+			--grpc_python_out=. \
+			./protos/$$v/*.proto ; \
+	done
+	$(call log_message,[INFO] gRPC Compilation complete!)
 
 grpc-server-start:
 	$(call log_message,INFO - Starting gRPC server ...)

grpc_publisher_service.py

@@ -1,36 +1,34 @@
 """gRPC Publisher Service"""
 
-import datetime
 import base64
-import traceback
+import datetime
 import json
-import grpc
+import traceback
 
+import grpc
 import sentry_sdk
 
-import publisher_pb2
-import publisher_pb2_grpc
-
-from utils import get_configs
 from content_parser import (
     decode_content,
     extract_content_v0,
     extract_content_v1,
     extract_content_v2,
 )
-from grpc_vault_entity_client import (
+from logutils import get_logger
+from notification_dispatcher import dispatch_notifications
+from platforms.adapter_ipc_handler import AdapterIPCHandler
+from platforms.adapter_manager import AdapterManager
+from protos.v1 import publisher_pb2, publisher_pb2_grpc
+from translations import Localization
+from utils import get_configs
+from vault_clients.v1.grpc_client import (
+    delete_entity_token,
+    get_entity_access_token,
     list_entity_stored_tokens,
     store_entity_token,
-    get_entity_access_token,
-    decrypt_payload,
     update_entity_token,
-    delete_entity_token,
 )
-from notification_dispatcher import dispatch_notifications
-from logutils import get_logger
-from translations import Localization
-from platforms.adapter_manager import AdapterManager
-from platforms.adapter_ipc_handler import AdapterIPCHandler
+from vault_clients.v2.grpc_client import decrypt_payload
 
 MOCK_DELIVERY_SMS = (
     get_configs("MOCK_DELIVERY_SMS", default_value="true") or ""

grpc_server.py

@@ -5,13 +5,13 @@
 
 import grpc
 from grpc_interceptor import ServerInterceptor
-import publisher_pb2_grpc
 
-from utils import get_configs
-from logutils import get_logger
-from sentry_config import initialize_sentry, SENTRY_ENABLED
 from grpc_publisher_service import PublisherService
+from logutils import get_logger
 from platforms.adapter_manager import AdapterManager
+from protos.v1 import publisher_pb2_grpc
+from sentry_config import SENTRY_ENABLED, initialize_sentry
+from utils import get_configs
 
 logger = get_logger("publisher.grpc.server")
 

vault_clients/__init__.py

@@ -1,13 +1,12 @@
-"""Vault gRPC Client"""
+"""Vault gRPC Client V1"""
 
 import functools
-import grpc
 
-import vault_pb2
-import vault_pb2_grpc
+import grpc
 
-from utils import get_configs, mask_sensitive_info
 from logutils import get_logger
+from protos.v1 import vault_pb2, vault_pb2_grpc
+from utils import get_configs, mask_sensitive_info
 
 logger = get_logger(__name__)
 
@@ -159,7 +158,9 @@ def get_entity_access_token(platform, account_identifier, **kwargs):
     id_type = (
         "long_lived_token"
         if long_lived_token
-        else "device_id" if device_id else "phone_number"
+        else "device_id"
+        if device_id
+        else "phone_number"
     )
 
     logger.debug("Requesting access tokens using %s='%s'...", id_type, identifier)

vault_clients/v1/__init__.py

@@ -0,0 +1,107 @@
+"""Vault gRPC Client V2"""
+
+import functools
+
+import grpc
+
+from logutils import get_logger
+from protos.v2 import vault_pb2, vault_pb2_grpc
+from utils import get_configs, mask_sensitive_info
+
+logger = get_logger(__name__)
+
+
+def get_channel(internal=True):
+    """Get the appropriate gRPC channel based on the mode.
+
+    Args:
+        internal (bool, optional): Flag indicating whether to use internal ports.
+            Defaults to True.
+
+    Returns:
+        grpc.Channel: The gRPC channel.
+    """
+    mode = get_configs("MODE", default_value="development")
+    hostname = get_configs("VAULT_GRPC_HOST")
+    if internal:
+        port = get_configs("VAULT_GRPC_INTERNAL_PORT")
+        secure_port = get_configs("VAULT_GRPC_INTERNAL_SSL_PORT")
+    else:
+        port = get_configs("VAULT_GRPC_PORT")
+        secure_port = get_configs("VAULT_GRPC_SSL_PORT")
+
+    if mode == "production":
+        logger.info("Connecting to vault gRPC server at %s:%s", hostname, secure_port)
+        credentials = grpc.ssl_channel_credentials()
+        logger.info("Using secure channel for gRPC communication")
+        return grpc.secure_channel(f"{hostname}:{secure_port}", credentials)
+
+    logger.info("Connecting to vault gRPC server at %s:%s", hostname, port)
+    logger.warning("Using insecure channel for gRPC communication")
+    return grpc.insecure_channel(f"{hostname}:{port}")
+
+
+def grpc_call(internal=True):
+    """Decorator to handle gRPC calls."""
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                channel = get_channel(internal)
+
+                with channel as conn:
+                    kwargs["stub"] = (
+                        vault_pb2_grpc.EntityInternalStub(conn)
+                        if internal
+                        else vault_pb2_grpc.EntityStub(conn)
+                    )
+                    return func(*args, **kwargs)
+            except grpc.RpcError as e:
+                return None, e
+            except Exception as e:
+                raise e
+
+        return wrapper
+
+    return decorator
+
+
+@grpc_call()
+def decrypt_payload(payload_ciphertext, **kwargs):
+    """
+    Decrypts the payload.
+
+    Args:
+        payload_ciphertext (bytes): The ciphertext of the payload to be decrypted.
+
+    Returns:
+        tuple: A tuple containing:
+            - server response (object): The vault server response.
+            - error (Exception): The error encountered if the request fails, otherwise None.
+    """
+    stub = kwargs["stub"]
+    device_id = kwargs.get("device_id")
+    phone_number = kwargs.get("phone_number")
+
+    request = vault_pb2.DecryptPayloadRequest(
+        device_id=device_id,
+        payload_ciphertext=payload_ciphertext,
+        phone_number=phone_number,
+    )
+
+    identifier = mask_sensitive_info(device_id or phone_number)
+
+    logger.debug(
+        "Initiating decryption request using %s='%s'.",
+        "device_id" if device_id else "phone_number",
+        identifier,
+    )
+
+    response = stub.DecryptPayload(request)
+
+    logger.info(
+        "Decryption successful using %s.",
+        "device_id" if device_id else "phone_number",
+    )
+    return response, None