Merge pull request #111 from smswithoutborders/feature/grpc_api

feat: complete encrypt paylaod function

src/grpc_entity_internal_service.py

@@ -24,6 +24,8 @@
 from src.relaysms_payload import (
     decode_relay_sms_payload,
     decrypt_payload,
+    encrypt_payload,
+    encode_relay_sms_payload,
 )
 
 HASHING_KEY = load_key(get_configs("HASHING_SALT"), 32)
@@ -357,6 +359,102 @@ def decrypt_message(entity_obj, header, content_ciphertext):
                 _type="UNKNOWN",
             )
 
+    def EncryptPayload(self, request, context):
+        """Handles encrypting payload"""
+
+        response = vault_pb2.EncryptPayloadResponse
+
+        def validate_fields():
+            return validate_request_fields(
+                context,
+                request,
+                response,
+                ["device_id", "payload_plaintext"],
+            )
+
+        def encrypt_message(entity_obj):
+            header, content_ciphertext, state, encrypt_error = encrypt_payload(
+                server_state=entity_obj.server_state,
+                client_publish_pub_key=base64.b64decode(
+                    entity_obj.client_publish_pub_key
+                ),
+                content=request.payload_plaintext,
+            )
+
+            if encrypt_error:
+                return error_response(
+                    context,
+                    response,
+                    encrypt_error,
+                    grpc.StatusCode.INVALID_ARGUMENT,
+                    user_msg="Invalid content format.",
+                    _type="UNKNOWN",
+                )
+
+            return (header, content_ciphertext, state), None
+
+        def encode_message(header, content_ciphertext, state):
+            encoded_content, encode_error = encode_relay_sms_payload(
+                header, content_ciphertext
+            )
+
+            if encode_error:
+                return None, error_response(
+                    context,
+                    response,
+                    encode_error,
+                    grpc.StatusCode.INVALID_ARGUMENT,
+                    user_msg="Invalid content format.",
+                    _type="UNKNOWN",
+                )
+
+            entity_obj.server_state = state.serialize()
+            entity_obj.save()
+            logger.info(
+                "Successfully encrypted payload for %s",
+                entity_obj.eid,
+            )
+
+            return response(
+                message="Successfully encrypted payload.",
+                payload_ciphertext=encoded_content,
+                success=True,
+            )
+
+        try:
+            invalid_fields_response = validate_fields()
+            if invalid_fields_response:
+                return invalid_fields_response
+
+            entity_obj = find_entity(device_id=request.device_id)
+
+            if not entity_obj:
+                return error_response(
+                    context,
+                    response,
+                    f"Invalid device ID '{request.device_id}'. "
+                    "Please log in again to obtain a valid device ID.",
+                    grpc.StatusCode.UNAUTHENTICATED,
+                )
+
+            encrypted_response, encrypting_error = encrypt_message(entity_obj)
+            if encrypting_error:
+                return encrypting_error
+
+            header, content_ciphertext, state = encrypted_response
+
+            return encode_message(header, content_ciphertext, state)
+
+        except Exception as e:
+            return error_response(
+                context,
+                response,
+                e,
+                grpc.StatusCode.INTERNAL,
+                user_msg="Oops! Something went wrong. Please try again later.",
+                _type="UNKNOWN",
+            )
+
     def GetEntityAccessToken(self, request, context):
         """Handles getting an entity's access token."""
 
@@ -457,70 +555,6 @@ def fetch_tokens(entity_obj, account_identifier_hash):
                 _type="UNKNOWN",
             )
 
-    def EncryptPayload(self, request, context):
-        """Handles encrypting payload"""
-
-        response = vault_pb2.EncryptPayloadResponse
-
-        try:
-            invalid_fields_response = validate_request_fields(
-                context,
-                request,
-                response,
-                ["device_id", "payload_plaintext"],
-            )
-            if invalid_fields_response:
-                return invalid_fields_response
-
-            entity_obj = find_entity(device_id=request.device_id)
-
-            if not entity_obj:
-                return error_response(
-                    context,
-                    response,
-                    f"Invalid device ID '{request.device_id}'. "
-                    "Please log in again to obtain a valid device ID.",
-                    grpc.StatusCode.UNAUTHENTICATED,
-                )
-
-            # entity_publish_keypair = load_keypair_object(entity_obj.publish_keypair)
-            # publish_shared_key = entity_publish_keypair.agree(
-            #     base64.b64decode(entity_obj.client_publish_pub_key)
-            # )
-
-            # header, content_ciphertext, state = encrypt_payload(
-            #     server_state=entity_obj.server_state,
-            #     publish_shared_key=publish_shared_key,
-            #     keypair=load_keypair_object(entity_obj.publish_keypair),
-            #     content=request.payload_plaintext,
-            #     client_pub_key=base64.b64decode(entity_obj.client_publish_pub_key),
-            #     client_keystore_path=os.path.join(
-            #         KEYSTORE_PATH, f"{entity_obj.eid.hex}_publish.db"
-            #     ),
-            # )
-
-            # b64_encoded_content = encode_relay_sms_payload(header, content_ciphertext)
-
-            # entity_obj.server_state = state.serialize()
-            # entity_obj.save()
-            b64_encoded_content = request.payload_plaintext
-
-            return response(
-                message="Successfully encrypted payload.",
-                payload_ciphertext=b64_encoded_content,
-                success=True,
-            )
-
-        except Exception as e:
-            return error_response(
-                context,
-                response,
-                e,
-                grpc.StatusCode.INTERNAL,
-                user_msg="Oops! Something went wrong. Please try again later.",
-                _type="UNKNOWN",
-            )
-
     def UpdateEntityToken(self, request, context):
         """Handles updating tokens for an entity"""
 

src/relaysms_payload.py

@@ -58,37 +58,33 @@ def decrypt_payload(
         return None, None, e
 
 
-def encrypt_payload(
-    server_state, publish_shared_key, client_publish_pub_key, content, **kwargs
-):
+def encrypt_payload(server_state, client_publish_pub_key, content):
     """
     Encrypts content into a RelaySMS payload.
 
     Args:
         server_state (bytes): Current state of the server-side ratchet.
-        publish_shared_key (bytes): Publish shared key.
         client_publish_pub_key (bytes): Client's public key for encryption.
         content (str): Plaintext content to encrypt.
-        kwargs (dict): Additional keyword arguments:
-            - client_keystore_path (str): Path to client's keystore.
 
     Returns:
         tuple:
             - header (bytes): Serialized ratchet header.
             - content_ciphertext (bytes): Encrypted content.
             - state (bytes): Updated server state.
+            - error (Exception or None)
     """
-    state = States.deserialize(server_state)
-    server_keystore_path = kwargs.get("server_keystore_path")
+    try:
+        state = States.deserialize(server_state)
 
-    Ratchets.alice_init(
-        state, publish_shared_key, client_publish_pub_key, server_keystore_path
-    )
-    header, content_ciphertext = Ratchets.encrypt(
-        state, content, client_publish_pub_key
-    )
+        header, content_ciphertext = Ratchets.encrypt(
+            state, content.encode("utf-8"), client_publish_pub_key
+        )
 
-    return header.serialize(), content_ciphertext, state
+        return header.serialize(), content_ciphertext, state, None
+    except Exception as e:
+        logger.error("Error encrypting relaysms payload: %s", e, exc_info=True)
+        return None, None, None, e
 
 
 def decode_relay_sms_payload(content):
@@ -133,14 +129,19 @@ def encode_relay_sms_payload(header, content_ciphertext):
         content_ciphertext (bytes): Encrypted content.
 
     Returns:
-        str: Base64-encoded representation of the payload.
+        tuple:
+            - encrypted_payload (str): Base64-encoded representation of the payload.
+            - error (Exception or None)
     """
     try:
-        serialized_header = header.serialize()
-        len_header = len(serialized_header)
-        return base64.b64encode(
-            struct.pack("<i", len_header) + serialized_header + content_ciphertext
-        ).decode("utf-8")
-
-    except (struct.error, IndexError, base64.binascii.Error) as e:
-        raise ValueError("Invalid payload format") from e
+        len_header = len(header)
+        return (
+            base64.b64encode(
+                struct.pack("<i", len_header) + header + content_ciphertext
+            ).decode("utf-8"),
+            None,
+        )
+
+    except Exception as e:
+        logger.error("Error encoding relaysms payload: %s", e, exc_info=True)
+        return None, e