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: