From 3f0129a907b22c9cfa52f583e58e6cf6e5496c8a Mon Sep 17 00:00:00 2001
From: Promise Fru <info@promisefru.com>
Date: Thu, 16 Jan 2025 22:04:55 +0100
Subject: [PATCH] feat: add support for encrypting payloads using a phone
 number as an identifier, in addition to the existing device ID.

---
 protos/v1/vault.proto               |  2 ++
 src/grpc_entity_internal_service.py | 33 +++++++++++++++++++----------
 2 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/protos/v1/vault.proto b/protos/v1/vault.proto
index 6cef15b..1cbcabd 100644
--- a/protos/v1/vault.proto
+++ b/protos/v1/vault.proto
@@ -157,6 +157,8 @@ message EncryptPayloadRequest {
   string device_id = 1;
   // Plaintext payload to be encrypted.
   string payload_plaintext = 2;
+  // The phone number of the entity.
+  string phone_number = 3;
 }
 
 // Response message for encrypting payload.
diff --git a/src/grpc_entity_internal_service.py b/src/grpc_entity_internal_service.py
index c43b210..c35d0ea 100644
--- a/src/grpc_entity_internal_service.py
+++ b/src/grpc_entity_internal_service.py
@@ -457,7 +457,7 @@ def validate_fields():
                 context,
                 request,
                 response,
-                ["device_id", "payload_plaintext"],
+                [("device_id", "phone_number"), "payload_plaintext"],
             )
 
         def encrypt_message(entity_obj):
@@ -514,16 +514,27 @@ def encode_message(header, content_ciphertext, state):
             if invalid_fields_response:
                 return invalid_fields_response
 
-            entity_obj = find_entity(device_id=request.device_id)
-
-            if not entity_obj:
-                return self.handle_create_grpc_error_response(
-                    context,
-                    response,
-                    f"Invalid device ID '{request.device_id}'. "
-                    "Please log in again to obtain a valid device ID.",
-                    grpc.StatusCode.UNAUTHENTICATED,
-                )
+            if request.device_id:
+                entity_obj = find_entity(device_id=request.device_id)
+                if not entity_obj:
+                    return self.handle_create_grpc_error_response(
+                        context,
+                        response,
+                        f"Entity associated with device ID '{request.device_id}' not found. "
+                        "Please log in again to obtain a valid device ID.",
+                        grpc.StatusCode.UNAUTHENTICATED,
+                    )
+            else:
+                phone_number_hash = generate_hmac(HASHING_KEY, request.phone_number)
+                entity_obj = find_entity(phone_number_hash=phone_number_hash)
+                if not entity_obj:
+                    return self.handle_create_grpc_error_response(
+                        context,
+                        response,
+                        f"Entity associated with phone number '{request.phone_number}' not found. "
+                        "Please check your phone number and try again.",
+                        grpc.StatusCode.UNAUTHENTICATED,
+                    )
 
             encrypted_response, encrypting_error = encrypt_message(entity_obj)
             if encrypting_error: