From 095852ce03b7da3e3932b09cae43f70ac1f53335 Mon Sep 17 00:00:00 2001
From: Promise Fru <info@promisefru.com>
Date: Sat, 15 Jun 2024 02:03:39 +0100
Subject: [PATCH 1/6] feat(api): basic implementation of the ratchet

---
 protos/v1/vault.proto      | 24 ++++++++++
 src/db_models.py           | 13 +++++-
 src/grpc_entity_service.py | 89 ++++++++++++++++++++++++++++++++++++++
 src/relaysms_payload.py    | 60 ++++++++++++++++++++++---
 tests/sample.py            | 83 +++++++++++++++++++++++++++++++++++
 5 files changed, 263 insertions(+), 6 deletions(-)
 create mode 100644 tests/sample.py

diff --git a/protos/v1/vault.proto b/protos/v1/vault.proto
index 7420f2a..06996c2 100644
--- a/protos/v1/vault.proto
+++ b/protos/v1/vault.proto
@@ -106,6 +106,28 @@ message StoreEntityTokenResponse {
     bool success = 2;
 }
 
+// Request message for getting entity access token and decrypting payload.
+message GetEntityAccessTokenAndDecryptPayloadRequest {
+    // Device ID for identifying the requesting device.
+    string device_id = 1;
+    // Encrypted payload that needs to be decrypted.
+    string payload_ciphertext = 2;
+}
+
+// Response message for getting entity access token and decrypting payload.
+message GetEntityAccessTokenAndDecryptPayloadResponse {
+    // Entity access token obtained from the operation.
+    string token = 1;
+    // Decrypted plaintext payload.
+    string payload_plaintext = 2;
+    // A response message.
+    string message = 3;
+    // Indicates whether the operation was successful.
+    bool success = 4;
+    // The platform associated with the token.
+    string platform = 5;
+}
+
 // Service for managing entities.
 service Entity {
     // Creates an entity.
@@ -116,4 +138,6 @@ service Entity {
     rpc ListEntityStoredTokens (ListEntityStoredTokenRequest) returns (ListEntityStoredTokenResponse);
     // Stores a token for an entity.
     rpc StoreEntityToken (StoreEntityTokenRequest) returns (StoreEntityTokenResponse);
+    // Get an entity's access token and decrypt payload.
+    rpc GetEntityAccessTokenAndDecryptPayload (GetEntityAccessTokenAndDecryptPayloadRequest) returns (GetEntityAccessTokenAndDecryptPayloadResponse);
 }
diff --git a/src/db_models.py b/src/db_models.py
index 016bc41..8d6a507 100644
--- a/src/db_models.py
+++ b/src/db_models.py
@@ -9,6 +9,7 @@
     IntegerField,
     UUIDField,
     ForeignKeyField,
+    BlobField,
 )
 from src.db import connect
 from src.utils import create_tables
@@ -28,6 +29,7 @@ class Entity(Model):
     client_publish_pub_key = TextField(null=True)
     client_device_id_pub_key = TextField(null=True)
     server_crypto_metadata = TextField(null=True)
+    server_state = BlobField(null=True)
     date_created = DateTimeField(default=datetime.datetime.now)
 
     class Meta:
@@ -35,7 +37,16 @@ class Meta:
 
         database = database
         table_name = "entities"
-        indexes = ((("phone_number_hash",), True),)
+        indexes = (
+            (
+                ("phone_number_hash",),
+                True,
+            ),
+            (
+                ("device_id",),
+                True,
+            ),
+        )
 
 
 class OTPRateLimit(Model):
diff --git a/src/grpc_entity_service.py b/src/grpc_entity_service.py
index 9c4fcca..7c3bad3 100644
--- a/src/grpc_entity_service.py
+++ b/src/grpc_entity_service.py
@@ -6,6 +6,8 @@
 import base64
 
 import grpc
+from smswithoutborders_libsig.ratchets import States
+from smswithoutborders_libsig.keypairs import x25519
 
 import vault_pb2
 import vault_pb2_grpc
@@ -28,6 +30,7 @@
 )
 from src.long_lived_token import generate_llt, verify_llt
 from src.device_id import compute_device_id
+from src.relaysms_payload import decode_relay_sms_payload, initialize_ratchet
 
 HASHING_KEY = load_key(get_configs("HASHING_SALT"), 32)
 KEYSTORE_PATH = get_configs("KEYSTORE_PATH")
@@ -284,6 +287,7 @@ def complete_creation(request):
                 "client_publish_pub_key": request.client_publish_pub_key,
                 "client_device_id_pub_key": request.client_device_id_pub_key,
                 "server_crypto_metadata": crypto_metadata_ciphertext_b64,
+                "server_state": States(),
             }
 
             create_entity(**fields)
@@ -480,6 +484,7 @@ def complete_authentication(request, entity_obj):
             entity_obj.client_publish_pub_key = request.client_publish_pub_key
             entity_obj.client_device_id_pub_key = request.client_device_id_pub_key
             entity_obj.server_crypto_metadata = crypto_metadata_ciphertext_b64
+            entity_obj.server_state = States()
             entity_obj.save()
 
             return response(
@@ -643,3 +648,87 @@ def StoreEntityToken(self, request, context):
                 user_msg="Oops! Something went wrong. Please try again later.",
                 _type="UNKNOWN",
             )
+
+    def GetEntityAccessTokenAndDecryptPayload(self, request, context):
+        """Handles getting an entity's access token and decrypting payload"""
+
+        response = vault_pb2.GetEntityAccessTokenAndDecryptPayloadResponse
+
+        try:
+            invalid_fields_response = validate_request_fields(
+                context,
+                request,
+                response,
+                ["device_id", "payload_ciphertext"],
+            )
+            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_crypto_metadata = load_crypto_metadata(
+                decrypt_and_decode(entity_obj.server_crypto_metadata)
+            )
+            entity_publish_keypair = entity_crypto_metadata.publish_keypair
+            entity_publish_keypair_obj = x25519(
+                os.path.join(KEYSTORE_PATH, f"{entity_obj.eid.hex}_publish.db"),
+                entity_publish_keypair.pnt_keystore,
+                entity_publish_keypair.secret_key,
+            )
+            entity_publish_shared_key = entity_publish_keypair_obj.agree(
+                base64.b64decode(entity_obj.client_publish_pub_key)
+            )
+
+            header, content_ciphertext = decode_relay_sms_payload(
+                request.payload_ciphertext
+            )
+
+            state = States.deserialize(entity_obj.server_state)
+
+            content_plaintext = initialize_ratchet(
+                state,
+                entity_publish_shared_key,
+                entity_publish_keypair_obj,
+                header,
+                content_ciphertext,
+                base64.b64decode(entity_obj.client_publish_pub_key),
+            )
+
+            tokens = fetch_entity_tokens(
+                entity=entity_obj,
+                fields=["account_tokens"],
+                platform=request.platform,
+            )
+            for token in tokens:
+                for field in ["account_tokens"]:
+                    if field in token:
+                        token[field] = decrypt_and_decode(token[field])
+
+            entity_obj.server_state = state.serialize()
+            entity_obj.save()
+
+            return response(
+                message="Successfully fetched tokens and decrypted payload",
+                success=True,
+                payload_plaintext=content_plaintext,
+                token=json.loads(tokens[0]),
+            )
+
+        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",
+            )
diff --git a/src/relaysms_payload.py b/src/relaysms_payload.py
index a391e19..d552b35 100644
--- a/src/relaysms_payload.py
+++ b/src/relaysms_payload.py
@@ -1,11 +1,61 @@
-""""""
+"""Module for decoding and extracting information from RelaySMS payloads."""
 
+import base64
+import struct
 from smswithoutborders_libsig.ratchets import Ratchets
 
 
-def destructure_paylaod(payload):
-    """"""
+def initialize_ratchet(
+    state, share_key, keypair_obj, header, ciphertext, client_public_key
+):
+    """
+    Initialize a ratchet.
 
+    Args:
+    - state (obj): State object for ratchet initialization.
+    - share_key (bytes): Shared key for cryptographic operations.
+    - keypair_obj (obj): Object containing keypair information.
+    - header (bytes): Encrypted header data to decrypt.
+    - ciphertext (bytes): Encrypted content data.
+    - client_public_key (bytes): Public key of the client.
 
-def decrypt_ciphertext(ciphertext):
-    """"""
+    Returns:
+    - bytes: Decrypted plaintext header.
+    """
+
+    Ratchets.bob_init(state, share_key, keypair_obj)
+    plaintext_header = Ratchets.decrypt(state, header, ciphertext, client_public_key)
+    return plaintext_header
+
+
+def decode_relay_sms_payload(content):
+    """
+    Decode a relay SMS payload containing a header and encrypted content.
+
+    Args:
+    - content (str): Base64-encoded string representing the payload.
+
+    Returns:
+    - tuple: A tuple containing the header (bytes) and encrypted content (bytes).
+
+    Raises:
+    - ValueError: If the payload format is invalid or decoding fails.
+    """
+    if not isinstance(content, str) or not content:
+        raise ValueError("Invalid input: content must be a non-empty string")
+
+    try:
+        # Decode base64 content
+        payload = base64.b64decode(content)
+
+        # Extract length of header
+        len_header = struct.unpack("<i", payload[:4])[0]
+
+        # Extract header and encrypted content
+        header = payload[4 : 4 + len_header]
+        encrypted_content = payload[4 + len_header :]
+
+        return header, encrypted_content
+
+    except (struct.error, IndexError, base64.binascii.Error) as e:
+        raise ValueError("Invalid payload format") from e
diff --git a/tests/sample.py b/tests/sample.py
new file mode 100644
index 0000000..cff6704
--- /dev/null
+++ b/tests/sample.py
@@ -0,0 +1,83 @@
+import base64
+import struct
+import grpc
+
+import vault_pb2
+import vault_pb2_grpc
+
+from smswithoutborders_libsig.keypairs import x25519
+from smswithoutborders_libsig.ratchets import Ratchets, States
+from src.device_id import compute_device_id
+
+
+def main():
+    try:
+        channel = grpc.insecure_channel("localhost:6000")
+
+        with channel as conn:
+            stub = vault_pb2_grpc.EntityStub(conn)
+
+            client_path = "/home/megamind/promisefru/sample.db"
+            client_path_1 = "/home/megamind/promisefru/sample1.db"
+            keypair_obj = x25519(client_path)
+            peer_pub_key = keypair_obj.init()
+
+            a_request_data = {
+                "phone_number": "+23712345678900",
+                "ownership_proof_response": "123456",
+                "country_code": "CM",
+                "password": "Password@1234",
+                "client_publish_pub_key": base64.b64encode(peer_pub_key).decode(
+                    "utf-8"
+                ),
+                "client_device_id_pub_key": base64.b64encode(peer_pub_key).decode(
+                    "utf-8"
+                ),
+            }
+
+            a_request = vault_pb2.CreateEntityRequest(**a_request_data)
+            a_response = stub.CreateEntity(a_request)
+
+            print(a_response)
+
+            server_pub = base64.b64decode(a_response.server_publish_pub_key)
+            sk = keypair_obj.agree(server_pub)
+
+            original_plaintext = b"Hello world"
+            client_state = States()
+            Ratchets.alice_init(
+                client_state,
+                sk,
+                server_pub,
+                client_path_1,
+            )
+            header, client_ciphertext = Ratchets.encrypt(
+                client_state,
+                original_plaintext,
+                server_pub,
+            )
+
+            len_header = len(header)
+            transmission_text = base64.b64encode(
+                struct.pack("<i", len_header) + header + client_ciphertext
+            )
+
+            b_request_data = {
+                "device_id": compute_device_id(sk, "+23712345678900", server_pub),
+                "payload_ciphertext": transmission_text,
+            }
+
+            b_request = vault_pb2.GetEntityAccessTokenAndDecryptPayloadRequest(
+                **b_request_data
+            )
+            b_response = stub.GetEntityAccessTokenAndDecryptPayload(b_request)
+
+            print(b_response)
+    except grpc.RpcError as e:
+        return None, e
+    except Exception as e:
+        raise e
+
+
+if __name__ == "__main__":
+    main()

From 951d979bc55ab7a8033defad6da0ff4ea0e88e00 Mon Sep 17 00:00:00 2001
From: Promise Fru <info@promisefru.com>
Date: Sun, 16 Jun 2024 18:57:07 +0100
Subject: [PATCH 2/6] feat: add update tokens to StoreEntityToken gRPC function

---
 protos/v1/vault.proto      |   6 +-
 src/db_models.py           |   3 +-
 src/device_id.py           |   2 +-
 src/grpc_entity_service.py | 213 ++++++++++++++++---------------------
 src/relaysms_payload.py    |  70 ++++++++----
 src/tokens.py              |  46 +++-----
 src/utils.py               | 117 ++++++++++----------
 7 files changed, 213 insertions(+), 244 deletions(-)

diff --git a/protos/v1/vault.proto b/protos/v1/vault.proto
index 06996c2..5e40a41 100644
--- a/protos/v1/vault.proto
+++ b/protos/v1/vault.proto
@@ -96,6 +96,8 @@ message StoreEntityTokenRequest {
     string platform = 3;
     // The identifier of the account associated with the token.
     string account_identifier = 4;
+    // Indicates whether to update existing tokens
+    bool update_token = 5;
 }
 
 // Response message for storing an entity's token.
@@ -112,6 +114,8 @@ message GetEntityAccessTokenAndDecryptPayloadRequest {
     string device_id = 1;
     // Encrypted payload that needs to be decrypted.
     string payload_ciphertext = 2;
+    // The platform associated with the token.
+    string platform = 3;
 }
 
 // Response message for getting entity access token and decrypting payload.
@@ -124,8 +128,6 @@ message GetEntityAccessTokenAndDecryptPayloadResponse {
     string message = 3;
     // Indicates whether the operation was successful.
     bool success = 4;
-    // The platform associated with the token.
-    string platform = 5;
 }
 
 // Service for managing entities.
diff --git a/src/db_models.py b/src/db_models.py
index 8d6a507..1bab7a6 100644
--- a/src/db_models.py
+++ b/src/db_models.py
@@ -28,7 +28,8 @@ class Entity(Model):
     device_id = CharField(null=True)
     client_publish_pub_key = TextField(null=True)
     client_device_id_pub_key = TextField(null=True)
-    server_crypto_metadata = TextField(null=True)
+    publish_keypair = BlobField(null=True)
+    device_id_keypair = BlobField(null=True)
     server_state = BlobField(null=True)
     date_created = DateTimeField(default=datetime.datetime.now)
 
diff --git a/src/device_id.py b/src/device_id.py
index cd75493..eca4410 100644
--- a/src/device_id.py
+++ b/src/device_id.py
@@ -19,5 +19,5 @@ def compute_device_id(secret_key, phone_number, public_key) -> str:
         str: The hexadecimal representation of the HMAC digest.
     """
     combined_input = phone_number + public_key
-    hmac_object = hmac.new(secret_key, combined_input.encode(), hashlib.sha256)
+    hmac_object = hmac.new(secret_key, combined_input.encode("utf-8"), hashlib.sha256)
     return hmac_object.hexdigest()
diff --git a/src/grpc_entity_service.py b/src/grpc_entity_service.py
index 7c3bad3..7c373ee 100644
--- a/src/grpc_entity_service.py
+++ b/src/grpc_entity_service.py
@@ -6,8 +6,6 @@
 import base64
 
 import grpc
-from smswithoutborders_libsig.ratchets import States
-from smswithoutborders_libsig.keypairs import x25519
 
 import vault_pb2
 import vault_pb2_grpc
@@ -21,12 +19,11 @@
     get_configs,
     encrypt_and_encode,
     generate_keypair_and_public_key,
-    generate_crypto_metadata,
     generate_eid,
-    get_shared_key,
     is_valid_x25519_public_key,
-    load_crypto_metadata,
     decrypt_and_decode,
+    load_keypair_object,
+    error_response,
 )
 from src.long_lived_token import generate_llt, verify_llt
 from src.device_id import compute_device_id
@@ -42,35 +39,6 @@
 logger = logging.getLogger("[gRPC Entity Service]")
 
 
-def error_response(context, response, sys_msg, status_code, user_msg=None, _type=None):
-    """
-    Create an error response.
-
-    Args:
-        context: gRPC context.
-        response: gRPC response object.
-        sys_msg (str or tuple): System message.
-        status_code: gRPC status code.
-        user_msg (str or tuple): User-friendly message.
-        _type (str): Type of error.
-
-    Returns:
-        An instance of the specified response with the error set.
-    """
-    if not user_msg:
-        user_msg = sys_msg
-
-    if _type == "UNKNOWN":
-        logger.exception(sys_msg, exc_info=True)
-    else:
-        logger.error(sys_msg)
-
-    context.set_details(user_msg)
-    context.set_code(status_code)
-
-    return response()
-
-
 def validate_request_fields(context, request, response, required_fields):
     """
     Validates the fields in the gRPC request.
@@ -201,14 +169,8 @@ def create_error_response(error_msg, error_detail=None):
             f"Possible token tampering detected. Entity not found with eid: {eid}"
         )
 
-    entity_crypto_metadata = load_crypto_metadata(
-        decrypt_and_decode(entity_obj.server_crypto_metadata)
-    )
-    entity_device_id_keypair = entity_crypto_metadata.device_id_keypair
-    entity_device_id_shared_key = get_shared_key(
-        os.path.join(KEYSTORE_PATH, f"{entity_obj.eid.hex}_device_id.db"),
-        entity_device_id_keypair.pnt_keystore,
-        entity_device_id_keypair.secret_key,
+    entity_device_id_keypair = load_keypair_object(entity_obj.device_id_keypair)
+    entity_device_id_shared_key = entity_device_id_keypair.agree(
         base64.b64decode(entity_obj.client_device_id_pub_key),
     )
 
@@ -251,30 +213,23 @@ def complete_creation(request):
             eid = generate_eid(phone_number_hash)
             password_hash = generate_hmac(HASHING_KEY, request.password)
             country_code_ciphertext_b64 = encrypt_and_encode(request.country_code)
-            server_publish_keypair, server_publish_pub_key = (
+
+            entity_publish_keypair, entity_publish_pub_key = (
                 generate_keypair_and_public_key(
                     os.path.join(KEYSTORE_PATH, f"{eid}_publish.db")
                 )
             )
-            server_device_id_keypair, server_device_id_pub_key = (
+            entity_device_id_keypair, entity_device_id_pub_key = (
                 generate_keypair_and_public_key(
                     os.path.join(KEYSTORE_PATH, f"{eid}_device_id.db")
                 )
             )
-            crypto_metadata_ciphertext_b64 = encrypt_and_encode(
-                generate_crypto_metadata(
-                    server_publish_keypair, server_device_id_keypair
-                )
-            )
 
-            shared_key = get_shared_key(
-                os.path.join(KEYSTORE_PATH, f"{eid}_device_id.db"),
-                server_device_id_keypair.pnt_keystore,
-                server_device_id_keypair.secret_key,
-                base64.b64decode(request.client_device_id_pub_key),
+            device_id_shared_key = entity_device_id_keypair.agree(
+                base64.b64decode(request.client_device_id_pub_key)
             )
 
-            long_lived_token = generate_llt(eid, shared_key)
+            long_lived_token = generate_llt(eid, device_id_shared_key)
 
             fields = {
                 "eid": eid,
@@ -282,12 +237,14 @@ def complete_creation(request):
                 "password_hash": password_hash,
                 "country_code": country_code_ciphertext_b64,
                 "device_id": compute_device_id(
-                    shared_key, request.phone_number, request.client_device_id_pub_key
+                    device_id_shared_key,
+                    request.phone_number,
+                    request.client_device_id_pub_key,
                 ),
                 "client_publish_pub_key": request.client_publish_pub_key,
                 "client_device_id_pub_key": request.client_device_id_pub_key,
-                "server_crypto_metadata": crypto_metadata_ciphertext_b64,
-                "server_state": States(),
+                "publish_keypair": entity_publish_keypair.serialize(),
+                "device_id_keypair": entity_device_id_keypair.serialize(),
             }
 
             create_entity(**fields)
@@ -297,11 +254,11 @@ def complete_creation(request):
             return response(
                 long_lived_token=long_lived_token,
                 message="Entity created successfully",
-                server_publish_pub_key=base64.b64encode(server_publish_pub_key).decode(
+                server_publish_pub_key=base64.b64encode(entity_publish_pub_key).decode(
                     "utf-8"
                 ),
                 server_device_id_pub_key=base64.b64encode(
-                    server_device_id_pub_key
+                    entity_device_id_pub_key
                 ).decode("utf-8"),
             )
 
@@ -453,48 +410,42 @@ def complete_authentication(request, entity_obj):
 
             eid = entity_obj.eid.hex
 
-            server_publish_keypair, server_publish_pub_key = (
+            entity_publish_keypair, entity_publish_pub_key = (
                 generate_keypair_and_public_key(
                     os.path.join(KEYSTORE_PATH, f"{eid}_publish.db")
                 )
             )
-            server_device_id_keypair, server_device_id_pub_key = (
+            entity_device_id_keypair, entity_device_id_pub_key = (
                 generate_keypair_and_public_key(
                     os.path.join(KEYSTORE_PATH, f"{eid}_device_id.db")
                 )
             )
-            crypto_metadata_ciphertext_b64 = encrypt_and_encode(
-                generate_crypto_metadata(
-                    server_publish_keypair, server_device_id_keypair
-                )
-            )
 
-            shared_key = get_shared_key(
-                os.path.join(KEYSTORE_PATH, f"{eid}_device_id.db"),
-                server_device_id_keypair.pnt_keystore,
-                server_device_id_keypair.secret_key,
-                base64.b64decode(request.client_device_id_pub_key),
+            device_id_shared_key = entity_device_id_keypair.agree(
+                base64.b64decode(request.client_device_id_pub_key)
             )
 
-            long_lived_token = generate_llt(eid, shared_key)
+            long_lived_token = generate_llt(eid, device_id_shared_key)
 
             entity_obj.device_id = compute_device_id(
-                shared_key, request.phone_number, request.client_device_id_pub_key
+                device_id_shared_key,
+                request.phone_number,
+                request.client_device_id_pub_key,
             )
             entity_obj.client_publish_pub_key = request.client_publish_pub_key
             entity_obj.client_device_id_pub_key = request.client_device_id_pub_key
-            entity_obj.server_crypto_metadata = crypto_metadata_ciphertext_b64
-            entity_obj.server_state = States()
+            entity_obj.publish_keypair = entity_publish_keypair.serialize()
+            entity_obj.device_id_keypair = entity_device_id_keypair.serialize()
             entity_obj.save()
 
             return response(
                 long_lived_token=long_lived_token,
                 message="Entity authenticated successfully!",
-                server_publish_pub_key=base64.b64encode(server_publish_pub_key).decode(
+                server_publish_pub_key=base64.b64encode(entity_publish_pub_key).decode(
                     "utf-8"
                 ),
                 server_device_id_pub_key=base64.b64encode(
-                    server_device_id_pub_key
+                    entity_device_id_pub_key
                 ).decode("utf-8"),
             )
 
@@ -550,7 +501,10 @@ def ListEntityStoredTokens(self, request, context):
                 return llt_error_response
 
             tokens = fetch_entity_tokens(
-                entity_obj, True, ["account_identifier", "platform"]
+                entity=entity_obj,
+                fetch_all=True,
+                fields=["account_identifier", "platform"],
+                return_json=True,
             )
             for token in tokens:
                 for field in ["account_identifier"]:
@@ -558,7 +512,6 @@ def ListEntityStoredTokens(self, request, context):
                         token[field] = decrypt_and_decode(token[field])
 
             logger.info("Successfully retrieved tokens for %s", entity_obj.eid)
-
             return response(
                 stored_tokens=tokens, message="Tokens retrieved successfully."
             )
@@ -601,35 +554,54 @@ def StoreEntityToken(self, request, context):
                     "this platform will be implemented."
                 )
 
-            if fetch_entity_tokens(
+            account_identifier_hash = generate_hmac(
+                HASHING_KEY, request.account_identifier
+            )
+
+            existing_tokens = fetch_entity_tokens(
                 entity=entity_obj,
-                account_identifier_hash=generate_hmac(
-                    HASHING_KEY, request.account_identifier
-                ),
+                account_identifier_hash=account_identifier_hash,
                 platform=request.platform,
-            ):
-                return error_response(
-                    context,
-                    response,
-                    "Entity already has a token associated with account "
-                    f"identifier {request.account_identifier} for {request.platform}",
-                    grpc.StatusCode.ALREADY_EXISTS,
-                )
-
-            new_token = {
-                "entity": entity_obj,
-                "platform": request.platform,
-                "account_identifier_hash": generate_hmac(
-                    HASHING_KEY, request.account_identifier
-                ),
-                "account_identifier": encrypt_and_encode(request.account_identifier),
-                "account_tokens": encrypt_and_encode(request.token),
-            }
-            create_entity_token(**new_token)
+            )
 
-            logger.info("Successfully stored tokens for %s", entity_obj.eid)
+            if request.update_token:
+                if not existing_tokens:
+                    return error_response(
+                        context,
+                        response,
+                        "No token found with account "
+                        f"identifier {request.account_identifier} for {request.platform}",
+                        grpc.StatusCode.NOT_FOUND,
+                    )
+
+                existing_tokens[0].account_tokens = encrypt_and_encode(request.token)
+                existing_tokens[0].save()
+                logger.info("Successfully updated token for %s", entity_obj.eid)
+            else:
+                if existing_tokens:
+                    return error_response(
+                        context,
+                        response,
+                        "Entity already has a token associated with account "
+                        f"identifier {request.account_identifier} for {request.platform}",
+                        grpc.StatusCode.ALREADY_EXISTS,
+                    )
+                new_token = {
+                    "entity": entity_obj,
+                    "platform": request.platform,
+                    "account_identifier_hash": generate_hmac(
+                        HASHING_KEY, request.account_identifier
+                    ),
+                    "account_identifier": encrypt_and_encode(
+                        request.account_identifier
+                    ),
+                    "account_tokens": encrypt_and_encode(request.token),
+                }
+                create_entity_token(**new_token)
+                logger.info("Successfully stored tokens for %s", entity_obj.eid)
 
-            return response(message="Token stored successfully.", success=True)
+            action = "updated" if request.update_token else "stored"
+            return response(message=f"Token {action} successfully.", success=True)
 
         except NotImplementedError as e:
             return error_response(
@@ -675,37 +647,30 @@ def GetEntityAccessTokenAndDecryptPayload(self, request, context):
                     grpc.StatusCode.UNAUTHENTICATED,
                 )
 
-            entity_crypto_metadata = load_crypto_metadata(
-                decrypt_and_decode(entity_obj.server_crypto_metadata)
-            )
-            entity_publish_keypair = entity_crypto_metadata.publish_keypair
-            entity_publish_keypair_obj = x25519(
-                os.path.join(KEYSTORE_PATH, f"{entity_obj.eid.hex}_publish.db"),
-                entity_publish_keypair.pnt_keystore,
-                entity_publish_keypair.secret_key,
-            )
-            entity_publish_shared_key = entity_publish_keypair_obj.agree(
+            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)
             )
 
+            print("publish_shared_key", publish_shared_key)
+
             header, content_ciphertext = decode_relay_sms_payload(
                 request.payload_ciphertext
             )
 
-            state = States.deserialize(entity_obj.server_state)
-
-            content_plaintext = initialize_ratchet(
-                state,
-                entity_publish_shared_key,
-                entity_publish_keypair_obj,
-                header,
-                content_ciphertext,
-                base64.b64decode(entity_obj.client_publish_pub_key),
+            content_plaintext, state = initialize_ratchet(
+                server_state=entity_obj.server_state,
+                publish_shared_key=publish_shared_key,
+                keypair=load_keypair_object(entity_obj.publish_keypair),
+                ratchet_header=header,
+                encrypted_content=content_ciphertext,
+                client_pub_key=base64.b64decode(entity_obj.client_publish_pub_key),
             )
 
             tokens = fetch_entity_tokens(
                 entity=entity_obj,
                 fields=["account_tokens"],
+                return_json=True,
                 platform=request.platform,
             )
             for token in tokens:
diff --git a/src/relaysms_payload.py b/src/relaysms_payload.py
index d552b35..664fc25 100644
--- a/src/relaysms_payload.py
+++ b/src/relaysms_payload.py
@@ -1,31 +1,58 @@
 """Module for decoding and extracting information from RelaySMS payloads."""
 
+import logging
 import base64
 import struct
-from smswithoutborders_libsig.ratchets import Ratchets
+from smswithoutborders_libsig.ratchets import Ratchets, States, HEADERS
+
+logging.basicConfig(
+    level=logging.INFO, format=("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+)
+logger = logging.getLogger(__name__)
 
 
 def initialize_ratchet(
-    state, share_key, keypair_obj, header, ciphertext, client_public_key
+    server_state, keypair, ratchet_header, encrypted_content, **kwargs
 ):
     """
     Initialize a ratchet.
 
     Args:
-    - state (obj): State object for ratchet initialization.
-    - share_key (bytes): Shared key for cryptographic operations.
-    - keypair_obj (obj): Object containing keypair information.
-    - header (bytes): Encrypted header data to decrypt.
-    - ciphertext (bytes): Encrypted content data.
-    - client_public_key (bytes): Public key of the client.
+        server_state (bytes or None): Current state of the server-side
+            ratchet. If None, initializes a new state.
+        keypair (object): Key pair object containing keys for encryption
+            and decryption.
+        ratchet_header (bytes): The ratchet header containing metadata for
+            the ratchet protocol.
+        encrypted_content (bytes): Encrypted content to be decrypted.
+        kwargs (dict): Additional keyword arguments:
+            - publish_shared_key (bytes): The shared key object to publish.
+            - client_pub_key (bytes): The client's public key for decryption.
 
     Returns:
-    - bytes: Decrypted plaintext header.
+        tuple:
+            - plaintext (str): Decrypted plaintext content.
+            - state (bytes): Updated server state.
     """
+    if not server_state:
+        state = States()
+    else:
+        state = States.deserialize(server_state)
+
+    publish_shared_key = kwargs.get("publish_shared_key")
+    client_pub_key = kwargs.get("client_pub_key")
+
+    Ratchets.bob_init(state, publish_shared_key, keypair)
+    logger.info("Ratchet initialized successfully.")
+
+    header = HEADERS(keypair)
+    header.deserialize(ratchet_header)
+    logger.info("Header deserialized successfully.")
 
-    Ratchets.bob_init(state, share_key, keypair_obj)
-    plaintext_header = Ratchets.decrypt(state, header, ciphertext, client_public_key)
-    return plaintext_header
+    plaintext = Ratchets.decrypt(state, header, encrypted_content, client_pub_key)
+    logger.info("Content decrypted successfully.")
+
+    return plaintext, state
 
 
 def decode_relay_sms_payload(content):
@@ -36,24 +63,23 @@ def decode_relay_sms_payload(content):
     - content (str): Base64-encoded string representing the payload.
 
     Returns:
-    - tuple: A tuple containing the header (bytes) and encrypted content (bytes).
-
-    Raises:
-    - ValueError: If the payload format is invalid or decoding fails.
+    - tuple:
+        - header (bytes): The ratchet header containing metadata for
+            the ratchet protocol.
+        - encrypted content (bytes): The encrypted payload.
     """
-    if not isinstance(content, str) or not content:
-        raise ValueError("Invalid input: content must be a non-empty string")
-
     try:
-        # Decode base64 content
         payload = base64.b64decode(content)
 
-        # Extract length of header
+        # Unpack the length of the header (first 4 bytes)
         len_header = struct.unpack("<i", payload[:4])[0]
 
-        # Extract header and encrypted content
+        # Extract the header (next len_header bytes)
         header = payload[4 : 4 + len_header]
+
+        # Extract the remaining payload as the encrypted content
         encrypted_content = payload[4 + len_header :]
+        logger.info("Header and encrypted content extracted.")
 
         return header, encrypted_content
 
diff --git a/src/tokens.py b/src/tokens.py
index 2fac4fe..5bd5fd1 100644
--- a/src/tokens.py
+++ b/src/tokens.py
@@ -4,6 +4,7 @@
 
 from playhouse.shortcuts import model_to_dict
 from src.db_models import Token
+from src.utils import remove_none_values
 
 database = Token._meta.database
 
@@ -43,23 +44,26 @@ def create_entity_token(
     return token
 
 
-def fetch_entity_tokens(entity, fetch_all=False, fields=None, **search_criteria):
+def fetch_entity_tokens(
+    entity, fetch_all=False, fields=None, return_json=False, **search_criteria
+):
     """
-    Fetch tokens associated with the given entity based on 
-    the provided search criteria, or retrieve all tokens 
-    associated with the entity if fetch_all is set to True.
+    Fetch tokens associated with the given entity.
 
     Args:
         entity (Entity): The entity associated with the tokens.
-        fetch_all (bool, optional): If True, fetch all tokens 
-            associated with the entity regardless of search criteria. 
+        fetch_all (bool, optional): If True, fetch all tokens
+            associated with the entity regardless of search criteria.
             If False (default), fetch tokens based on search criteria.
         fields (list[str] or None, optional): Optional list of fields to select.
+        return_json (bool, optional): If True, return the results as a list of dicts.
+            If False (default), return the results as a list of token objects.
         **search_criteria: Additional keyword arguments representing the fields
             and their values to search for.
 
     Returns:
-        list[dict]: A list of dictionaries containing token data.
+        list[dict] or list[Token]: A list of token objects or a list of dictionaries
+            containing token data if return_json is True.
     """
     results = []
 
@@ -78,32 +82,12 @@ def fetch_entity_tokens(entity, fetch_all=False, fields=None, **search_criteria)
             select = (getattr(Token, key) for key in fields)
             query = query.select(*select)
 
-        tokens = query.execute()
+        tokens = list(query.execute())
 
+    if return_json:
         for token in tokens:
             token_dict = model_to_dict(token)
             results.append(token_dict)
+        return remove_none_values(results)
 
-    return remove_none_values(results)
-
-
-def remove_none_values(values):
-    """
-    Removes None values from a list of dictionaries.
-
-    Args:
-        values (list): A list of dictionaries.
-
-    Returns:
-        list: A new list of dictionaries where None values have been removed.
-
-    Example:
-        values = [
-            {"a": 1, "b": None, "c": 3},
-            {"a": None, "b": 2, "c": 3},
-            {"a": 1, "b": 2, "c": None}
-        ]
-        filtered_values = remove_none_values(values)
-        # Output: [{'a': 1, 'c': 3}, {'b': 2, 'c': 3}, {'a': 1, 'b': 2}]
-    """
-    return [{k: v for k, v in value.items() if v is not None} for value in values]
+    return tokens
diff --git a/src/utils.py b/src/utils.py
index 232fbf1..30d9633 100644
--- a/src/utils.py
+++ b/src/utils.py
@@ -3,7 +3,6 @@
 import os
 import logging
 import uuid
-import json
 import base64
 from functools import wraps
 
@@ -231,6 +230,20 @@ def generate_keypair_and_public_key(keystore_path):
     return keypair_obj, peer_pub_key
 
 
+def load_keypair_object(keypair):
+    """
+    Deserialize a serialized x25519 keypair object from bytes.
+
+    Args:
+        keypair (bytes): Serialized x25519 keypair object.
+
+    Returns:
+        x25519: Deserialized x25519 keypair object.
+    """
+    keypair_obj = x25519()
+    return keypair_obj.deserialize(keypair)
+
+
 def get_shared_key(keystore_path, pnt_keystore, secret_key, peer_pub_key):
     """
     Generate a shared key.
@@ -249,77 +262,33 @@ def get_shared_key(keystore_path, pnt_keystore, secret_key, peer_pub_key):
     return shared_key
 
 
-def generate_crypto_metadata(publish_keypair, device_id_keypair):
-    """
-    Generate cryptographic metadata.
-
-    Args:
-        publish_keypair: Publish keypair object.
-        device_id_keypair: Device ID keypair object.
-
-    Returns:
-        str: JSON string representing cryptographic metadata.
-    """
-    crypto_metadata = {
-        "publish_keypair": {
-            "pnt_keystore": publish_keypair.pnt_keystore,
-            "secret_key": publish_keypair.secret_key,
-        },
-        "device_id_keypair": {
-            "pnt_keystore": device_id_keypair.pnt_keystore,
-            "secret_key": device_id_keypair.secret_key,
-        },
-    }
-
-    return json.dumps(crypto_metadata)
-
-
-def load_crypto_metadata(metadata):
+def error_response(context, response, sys_msg, status_code, user_msg=None, _type=None):
     """
-    Load cryptographic metadata.
+    Create an error response.
 
     Args:
-        metadata (str): JSON string representing cryptographic metadata.
+        context: gRPC context.
+        response: gRPC response object.
+        sys_msg (str or tuple): System message.
+        status_code: gRPC status code.
+        user_msg (str or tuple): User-friendly message.
+        _type (str): Type of error.
 
     Returns:
-        CryptoMetadata: Object representing cryptographic metadata.
+        An instance of the specified response with the error set.
     """
+    if not user_msg:
+        user_msg = sys_msg
 
-    class CryptoMetadata:
-        """
-        Represents cryptographic metadata.
-        """
-
-        def __init__(self, metadata):
-            self.__dict__.update(self._convert(json.loads(metadata)))
-
-        def _convert(self, data):
-            for key, value in data.items():
-                if isinstance(value, dict):
-                    data[key] = CryptoMetadata(json.dumps(value))
-                elif isinstance(value, list):
-                    data[key] = [
-                        (
-                            CryptoMetadata(json.dumps(item))
-                            if isinstance(item, dict)
-                            else item
-                        )
-                        for item in value
-                    ]
-            return data
-
-        def __getattr__(self, name):
-            try:
-                return self.__dict__[name]
-            except KeyError as exc:
-                raise AttributeError(
-                    f"'CryptoMetadata' object has no attribute '{name}'"
-                ) from exc
+    if _type == "UNKNOWN":
+        logger.exception(sys_msg, exc_info=True)
+    else:
+        logger.error(sys_msg)
 
-        def __setattr__(self, name, value):
-            self.__dict__[name] = value
+    context.set_details(user_msg)
+    context.set_code(status_code)
 
-    return CryptoMetadata(metadata)
+    return response()
 
 
 def encrypt_and_encode(plaintext):
@@ -399,3 +368,25 @@ def is_valid_x25519_public_key(encoded_key):
     except ValueError as err:
         logger.exception("X25519 public key validation error: %s", err)
         return False, f"Invalid X25519 public key: {err}"
+
+
+def remove_none_values(values):
+    """
+    Removes None values from a list of dictionaries.
+
+    Args:
+        values (list): A list of dictionaries.
+
+    Returns:
+        list: A new list of dictionaries where None values have been removed.
+
+    Example:
+        values = [
+            {"a": 1, "b": None, "c": 3},
+            {"a": None, "b": 2, "c": 3},
+            {"a": 1, "b": 2, "c": None}
+        ]
+        filtered_values = remove_none_values(values)
+        # Output: [{'a': 1, 'c': 3}, {'b': 2, 'c': 3}, {'a': 1, 'b': 2}]
+    """
+    return [{k: v for k, v in value.items() if v is not None} for value in values]

From adf203eca20d7726c586e6e99af1cc058323d9cf Mon Sep 17 00:00:00 2001
From: Promise Fru <info@promisefru.com>
Date: Mon, 17 Jun 2024 00:18:33 +0100
Subject: [PATCH 3/6] feat: add EncryptPayload grpc function and completed
 relaysms payload module

---
 protos/v1/vault.proto      |  26 ++++++--
 src/grpc_entity_service.py | 133 +++++++++++++++++++++++++------------
 src/relaysms_payload.py    |  88 ++++++++++++++++++------
 3 files changed, 182 insertions(+), 65 deletions(-)

diff --git a/protos/v1/vault.proto b/protos/v1/vault.proto
index 5e40a41..91b4351 100644
--- a/protos/v1/vault.proto
+++ b/protos/v1/vault.proto
@@ -90,14 +90,12 @@ message Token {
 message StoreEntityTokenRequest {
     // The long-lived token of the authenticated entity.
     string long_lived_token = 1;
-    // The OAuth2 token to be stored.
+    // The OAuth2 token to be stored (JSON string).
     string token = 2;
     // The platform associated with the token.
     string platform = 3;
     // The identifier of the account associated with the token.
     string account_identifier = 4;
-    // Indicates whether to update existing tokens
-    bool update_token = 5;
 }
 
 // Response message for storing an entity's token.
@@ -120,7 +118,7 @@ message GetEntityAccessTokenAndDecryptPayloadRequest {
 
 // Response message for getting entity access token and decrypting payload.
 message GetEntityAccessTokenAndDecryptPayloadResponse {
-    // Entity access token obtained from the operation.
+    // Entity access token (JSON string).
     string token = 1;
     // Decrypted plaintext payload.
     string payload_plaintext = 2;
@@ -130,6 +128,24 @@ message GetEntityAccessTokenAndDecryptPayloadResponse {
     bool success = 4;
 }
 
+// Request message for encrypting payload.
+message EncryptPayloadRequest {
+    // Device ID for identifying the requesting device.
+    string device_id = 1;
+    // Plaintext payload to be encrypted.
+    string payload_plaintext = 2;
+}
+
+// Response message for encrypting payload.
+message EncryptPayloadResponse {
+    // Encrypted payload.
+    string payload_ciphertext = 1;
+    // A response message.
+    string message = 2;
+    // Indicates whether the operation was successful.
+    bool success = 3;
+}
+
 // Service for managing entities.
 service Entity {
     // Creates an entity.
@@ -142,4 +158,6 @@ service Entity {
     rpc StoreEntityToken (StoreEntityTokenRequest) returns (StoreEntityTokenResponse);
     // Get an entity's access token and decrypt payload.
     rpc GetEntityAccessTokenAndDecryptPayload (GetEntityAccessTokenAndDecryptPayloadRequest) returns (GetEntityAccessTokenAndDecryptPayloadResponse);
+    // Encrypt payload.
+    rpc EncryptPayload (EncryptPayloadRequest) returns (EncryptPayloadResponse);
 }
diff --git a/src/grpc_entity_service.py b/src/grpc_entity_service.py
index 7c373ee..b775bd2 100644
--- a/src/grpc_entity_service.py
+++ b/src/grpc_entity_service.py
@@ -27,7 +27,12 @@
 )
 from src.long_lived_token import generate_llt, verify_llt
 from src.device_id import compute_device_id
-from src.relaysms_payload import decode_relay_sms_payload, initialize_ratchet
+from src.relaysms_payload import (
+    decode_relay_sms_payload,
+    encode_relay_sms_payload,
+    encrypt_payload,
+    decrypt_payload,
+)
 
 HASHING_KEY = load_key(get_configs("HASHING_SALT"), 32)
 KEYSTORE_PATH = get_configs("KEYSTORE_PATH")
@@ -564,44 +569,29 @@ def StoreEntityToken(self, request, context):
                 platform=request.platform,
             )
 
-            if request.update_token:
-                if not existing_tokens:
-                    return error_response(
-                        context,
-                        response,
-                        "No token found with account "
-                        f"identifier {request.account_identifier} for {request.platform}",
-                        grpc.StatusCode.NOT_FOUND,
-                    )
-
-                existing_tokens[0].account_tokens = encrypt_and_encode(request.token)
-                existing_tokens[0].save()
-                logger.info("Successfully updated token for %s", entity_obj.eid)
-            else:
-                if existing_tokens:
-                    return error_response(
-                        context,
-                        response,
-                        "Entity already has a token associated with account "
-                        f"identifier {request.account_identifier} for {request.platform}",
-                        grpc.StatusCode.ALREADY_EXISTS,
-                    )
-                new_token = {
-                    "entity": entity_obj,
-                    "platform": request.platform,
-                    "account_identifier_hash": generate_hmac(
-                        HASHING_KEY, request.account_identifier
-                    ),
-                    "account_identifier": encrypt_and_encode(
-                        request.account_identifier
-                    ),
-                    "account_tokens": encrypt_and_encode(request.token),
-                }
-                create_entity_token(**new_token)
-                logger.info("Successfully stored tokens for %s", entity_obj.eid)
+            if existing_tokens:
+                return error_response(
+                    context,
+                    response,
+                    "Entity already has a token associated with account "
+                    f"identifier {request.account_identifier} for {request.platform}",
+                    grpc.StatusCode.ALREADY_EXISTS,
+                )
 
-            action = "updated" if request.update_token else "stored"
-            return response(message=f"Token {action} successfully.", success=True)
+            new_token = {
+                "entity": entity_obj,
+                "platform": request.platform,
+                "account_identifier_hash": account_identifier_hash,
+                "account_identifier": encrypt_and_encode(request.account_identifier),
+                "account_tokens": encrypt_and_encode(request.token),
+            }
+            create_entity_token(**new_token)
+            logger.info("Successfully stored tokens for %s", entity_obj.eid)
+
+            return response(
+                message=f"Token stored successfully.",
+                success=True,
+            )
 
         except NotImplementedError as e:
             return error_response(
@@ -652,13 +642,11 @@ def GetEntityAccessTokenAndDecryptPayload(self, request, context):
                 base64.b64decode(entity_obj.client_publish_pub_key)
             )
 
-            print("publish_shared_key", publish_shared_key)
-
             header, content_ciphertext = decode_relay_sms_payload(
                 request.payload_ciphertext
             )
 
-            content_plaintext, state = initialize_ratchet(
+            content_plaintext, state = decrypt_payload(
                 server_state=entity_obj.server_state,
                 publish_shared_key=publish_shared_key,
                 keypair=load_keypair_object(entity_obj.publish_keypair),
@@ -697,3 +685,66 @@ def GetEntityAccessTokenAndDecryptPayload(self, request, context):
                 user_msg="Oops! Something went wrong. Please try again later.",
                 _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()
+
+            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",
+            )
diff --git a/src/relaysms_payload.py b/src/relaysms_payload.py
index 664fc25..a9c0047 100644
--- a/src/relaysms_payload.py
+++ b/src/relaysms_payload.py
@@ -1,4 +1,6 @@
-"""Module for decoding and extracting information from RelaySMS payloads."""
+"""
+Module for handling encryption, decryption, encoding, and decoding of RelaySMS payloads.
+"""
 
 import logging
 import base64
@@ -6,28 +8,24 @@
 from smswithoutborders_libsig.ratchets import Ratchets, States, HEADERS
 
 logging.basicConfig(
-    level=logging.INFO, format=("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 logger = logging.getLogger(__name__)
 
 
-def initialize_ratchet(
-    server_state, keypair, ratchet_header, encrypted_content, **kwargs
-):
+def decrypt_payload(server_state, keypair, ratchet_header, encrypted_content, **kwargs):
     """
-    Initialize a ratchet.
+    Decrypts a RelaySMS payload.
 
     Args:
-        server_state (bytes or None): Current state of the server-side
-            ratchet. If None, initializes a new state.
-        keypair (object): Key pair object containing keys for encryption
-            and decryption.
-        ratchet_header (bytes): The ratchet header containing metadata for
-            the ratchet protocol.
+        server_state (bytes or None): Current state of the server-side ratchet.
+            If None, initializes a new state.
+        keypair (object): Object containing encryption and decryption keys.
+        ratchet_header (bytes): Ratchet header.
         encrypted_content (bytes): Encrypted content to be decrypted.
         kwargs (dict): Additional keyword arguments:
-            - publish_shared_key (bytes): The shared key object to publish.
-            - client_pub_key (bytes): The client's public key for decryption.
+            - publish_shared_key (bytes): Publish shared key.
+            - client_pub_key (bytes): Client's public key for decryption.
 
     Returns:
         tuple:
@@ -55,18 +53,46 @@ def initialize_ratchet(
     return plaintext, state
 
 
+def encrypt_payload(
+    server_state, publish_shared_key, client_pub_key, content, **kwargs
+):
+    """
+    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_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.
+    """
+    state = States.deserialize(server_state)
+    client_keystore_path = kwargs.get("client_keystore_path")
+
+    Ratchets.alice_init(state, publish_shared_key, client_pub_key, client_keystore_path)
+    header, content_ciphertext = Ratchets.encrypt(state, content, client_pub_key)
+
+    return header.serialize(), content_ciphertext, state
+
+
 def decode_relay_sms_payload(content):
     """
-    Decode a relay SMS payload containing a header and encrypted content.
+    Decode a RelaySMS payload from a base64-encoded string.
 
     Args:
-    - content (str): Base64-encoded string representing the payload.
+        content (str): Base64-encoded string representing the payload.
 
     Returns:
-    - tuple:
-        - header (bytes): The ratchet header containing metadata for
-            the ratchet protocol.
-        - encrypted content (bytes): The encrypted payload.
+        tuple:
+            - header (bytes): Ratchet header.
+            - encrypted_content (bytes): Encrypted payload.
     """
     try:
         payload = base64.b64decode(content)
@@ -85,3 +111,25 @@ def decode_relay_sms_payload(content):
 
     except (struct.error, IndexError, base64.binascii.Error) as e:
         raise ValueError("Invalid payload format") from e
+
+
+def encode_relay_sms_payload(header, content_ciphertext):
+    """
+    Encode a RelaySMS payload to a base64-encoded string.
+
+    Args:
+        header (bytes): Ratchet header.
+        content_ciphertext (bytes): Encrypted content.
+
+    Returns:
+        str: Base64-encoded representation of the payload.
+    """
+    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

From 67b0ce9047f4e665ce98302886b49909258545a1 Mon Sep 17 00:00:00 2001
From: Promise Fru <info@promisefru.com>
Date: Mon, 17 Jun 2024 12:30:34 +0100
Subject: [PATCH 4/6] feat: add UpdateEntityToken grpc function

---
 protos/v1/vault.proto      |  22 ++++++
 src/grpc_entity_service.py | 154 +++++++++++++++++++++++++++++--------
 src/relaysms_payload.py    |  37 +++++----
 3 files changed, 167 insertions(+), 46 deletions(-)

diff --git a/protos/v1/vault.proto b/protos/v1/vault.proto
index 91b4351..bbf5104 100644
--- a/protos/v1/vault.proto
+++ b/protos/v1/vault.proto
@@ -146,6 +146,26 @@ message EncryptPayloadResponse {
     bool success = 3;
 }
 
+// Request message for updating an entity's token.
+message UpdateEntityTokenRequest {
+    // Device ID for identifying the requesting device.
+    string device_id = 1;
+    // The OAuth2 token to be stored (JSON string).
+    string token = 2;
+    // The platform associated with the token.
+    string platform = 3;
+    // The identifier of the account associated with the token.
+    string account_identifier = 4;
+}
+
+// Response message for updating an entity's token.
+message UpdateEntityTokenResponse {
+    // A response message.
+    string message = 1;
+    // Indicates whether the operation was successful.
+    bool success = 2;
+}
+
 // Service for managing entities.
 service Entity {
     // Creates an entity.
@@ -160,4 +180,6 @@ service Entity {
     rpc GetEntityAccessTokenAndDecryptPayload (GetEntityAccessTokenAndDecryptPayloadRequest) returns (GetEntityAccessTokenAndDecryptPayloadResponse);
     // Encrypt payload.
     rpc EncryptPayload (EncryptPayloadRequest) returns (EncryptPayloadResponse);
+    // Updates a token for an entity.
+    rpc UpdateEntityToken (UpdateEntityTokenRequest) returns (UpdateEntityTokenResponse);
 }
diff --git a/src/grpc_entity_service.py b/src/grpc_entity_service.py
index b775bd2..80fcc79 100644
--- a/src/grpc_entity_service.py
+++ b/src/grpc_entity_service.py
@@ -621,7 +621,7 @@ def GetEntityAccessTokenAndDecryptPayload(self, request, context):
                 context,
                 request,
                 response,
-                ["device_id", "payload_ciphertext"],
+                ["device_id", "payload_ciphertext", "platform"],
             )
             if invalid_fields_response:
                 return invalid_fields_response
@@ -642,38 +642,67 @@ def GetEntityAccessTokenAndDecryptPayload(self, request, context):
                 base64.b64decode(entity_obj.client_publish_pub_key)
             )
 
-            header, content_ciphertext = decode_relay_sms_payload(
+            header, content_ciphertext, decode_error = decode_relay_sms_payload(
                 request.payload_ciphertext
             )
 
-            content_plaintext, state = decrypt_payload(
-                server_state=entity_obj.server_state,
-                publish_shared_key=publish_shared_key,
-                keypair=load_keypair_object(entity_obj.publish_keypair),
-                ratchet_header=header,
-                encrypted_content=content_ciphertext,
-                client_pub_key=base64.b64decode(entity_obj.client_publish_pub_key),
-            )
+            if decode_error:
+                return error_response(
+                    context,
+                    response,
+                    decode_error,
+                    grpc.StatusCode.INVALID_ARGUMENT,
+                    user_msg="Invalid content format.",
+                    _type="UNKNOWN",
+                )
 
+            content_plaintext = content_ciphertext.decode("utf-8")
+            # content_plaintext, state, decrypt_error = decrypt_payload(
+            #     server_state=entity_obj.server_state,
+            #     publish_shared_key=publish_shared_key,
+            #     keypair=load_keypair_object(entity_obj.publish_keypair),
+            #     ratchet_header=header,
+            #     encrypted_content=content_ciphertext,
+            #     client_pub_key=base64.b64decode(entity_obj.client_publish_pub_key),
+            # )
+
+            # if decrypt_error:
+            #     return error_response(
+            #         context,
+            #         response,
+            #         decrypt_error,
+            #         grpc.StatusCode.INVALID_ARGUMENT,
+            #         user_msg="Invalid content format.",
+            #         _type="UNKNOWN",
+            #     )
+
+            account_identifier_hash = generate_hmac(
+                HASHING_KEY, content_plaintext.split(":")[0]
+            )
             tokens = fetch_entity_tokens(
                 entity=entity_obj,
                 fields=["account_tokens"],
                 return_json=True,
                 platform=request.platform,
+                account_identifier_hash=account_identifier_hash,
             )
             for token in tokens:
                 for field in ["account_tokens"]:
                     if field in token:
                         token[field] = decrypt_and_decode(token[field])
 
-            entity_obj.server_state = state.serialize()
-            entity_obj.save()
+            # entity_obj.server_state = state.serialize()
+            # entity_obj.save()
+            logger.info(
+                "Successfully fetched tokens and decrypted payload for %s",
+                entity_obj.eid,
+            )
 
             return response(
                 message="Successfully fetched tokens and decrypted payload",
                 success=True,
                 payload_plaintext=content_plaintext,
-                token=json.loads(tokens[0]),
+                token=tokens[0]["account_tokens"],
             )
 
         except Exception as e:
@@ -712,26 +741,27 @@ def EncryptPayload(self, request, context):
                     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)
-            )
+            # 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"
-                ),
-            )
+            # 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)
+            # b64_encoded_content = encode_relay_sms_payload(header, content_ciphertext)
 
-            entity_obj.server_state = state.serialize()
-            entity_obj.save()
+            # entity_obj.server_state = state.serialize()
+            # entity_obj.save()
+            b64_encoded_content = request.payload_plaintext
 
             return response(
                 message="Successfully encrypted payload.",
@@ -748,3 +778,67 @@ def EncryptPayload(self, request, context):
                 user_msg="Oops! Something went wrong. Please try again later.",
                 _type="UNKNOWN",
             )
+
+    def UpdateEntityToken(self, request, context):
+        """Handles updating tokens for an entity"""
+
+        response = vault_pb2.UpdateEntityTokenResponse
+
+        try:
+            invalid_fields_response = validate_request_fields(
+                context,
+                request,
+                response,
+                ["device_id", "token", "platform", "account_identifier"],
+            )
+            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,
+                )
+
+            account_identifier_hash = generate_hmac(
+                HASHING_KEY, request.account_identifier
+            )
+
+            existing_tokens = fetch_entity_tokens(
+                entity=entity_obj,
+                account_identifier_hash=account_identifier_hash,
+                platform=request.platform,
+            )
+
+            if not existing_tokens:
+                return error_response(
+                    context,
+                    response,
+                    "No token found with account "
+                    f"identifier {request.account_identifier} for {request.platform}",
+                    grpc.StatusCode.NOT_FOUND,
+                )
+
+            existing_tokens[0].account_tokens = encrypt_and_encode(request.token)
+            existing_tokens[0].save()
+            logger.info("Successfully updated token for %s", entity_obj.eid)
+
+            return response(
+                message="Token updated successfully.",
+                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",
+            )
diff --git a/src/relaysms_payload.py b/src/relaysms_payload.py
index a9c0047..d8c767b 100644
--- a/src/relaysms_payload.py
+++ b/src/relaysms_payload.py
@@ -31,26 +31,30 @@ def decrypt_payload(server_state, keypair, ratchet_header, encrypted_content, **
         tuple:
             - plaintext (str): Decrypted plaintext content.
             - state (bytes): Updated server state.
+            - error (Exception or None)
     """
-    if not server_state:
-        state = States()
-    else:
-        state = States.deserialize(server_state)
+    try:
+        if not server_state:
+            state = States()
+        else:
+            state = States.deserialize(server_state)
 
-    publish_shared_key = kwargs.get("publish_shared_key")
-    client_pub_key = kwargs.get("client_pub_key")
+        publish_shared_key = kwargs.get("publish_shared_key")
+        client_pub_key = kwargs.get("client_pub_key")
 
-    Ratchets.bob_init(state, publish_shared_key, keypair)
-    logger.info("Ratchet initialized successfully.")
+        Ratchets.bob_init(state, publish_shared_key, keypair)
+        logger.info("Ratchet initialized successfully.")
 
-    header = HEADERS(keypair)
-    header.deserialize(ratchet_header)
-    logger.info("Header deserialized successfully.")
+        header = HEADERS(keypair)
+        header.deserialize(ratchet_header)
+        logger.info("Header deserialized successfully.")
 
-    plaintext = Ratchets.decrypt(state, header, encrypted_content, client_pub_key)
-    logger.info("Content decrypted successfully.")
+        plaintext = Ratchets.decrypt(state, header, encrypted_content, client_pub_key)
+        logger.info("Content decrypted successfully.")
 
-    return plaintext, state
+        return plaintext, state, None
+    except (struct.error, IndexError, base64.binascii.Error) as e:
+        return None, None, e
 
 
 def encrypt_payload(
@@ -93,6 +97,7 @@ def decode_relay_sms_payload(content):
         tuple:
             - header (bytes): Ratchet header.
             - encrypted_content (bytes): Encrypted payload.
+            - error (Exception or None)
     """
     try:
         payload = base64.b64decode(content)
@@ -107,10 +112,10 @@ def decode_relay_sms_payload(content):
         encrypted_content = payload[4 + len_header :]
         logger.info("Header and encrypted content extracted.")
 
-        return header, encrypted_content
+        return header, encrypted_content, None
 
     except (struct.error, IndexError, base64.binascii.Error) as e:
-        raise ValueError("Invalid payload format") from e
+        return None, None, e
 
 
 def encode_relay_sms_payload(header, content_ciphertext):

From cbfbeebfb345dc3fea1536ed1ee98283aab7e8f9 Mon Sep 17 00:00:00 2001
From: Promise Fru <info@promisefru.com>
Date: Mon, 17 Jun 2024 14:51:32 +0100
Subject: [PATCH 5/6] docs: add references for
 GetEntityAccessTokenAndDecryptPayload, EncryptPayload and UpdateEntityToken
 grpc functions

---
 docs/grpc.md | 234 ++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 233 insertions(+), 1 deletion(-)

diff --git a/docs/grpc.md b/docs/grpc.md
index 52b92b2..2f81361 100644
--- a/docs/grpc.md
+++ b/docs/grpc.md
@@ -14,6 +14,9 @@
     - [Complete Authentication](#complete-authentication)
   - [List an Entity's Stored Tokens](#list-an-entitys-stored-tokens)
   - [Store an Entity's Token](#store-an-entitys-token)
+  - [Get Entity Access Token and Decrypt Payload](#get-entity-access-token-and-decrypt-payload)
+  - [Encrypt Payload](#encrypt-payload)
+  - [Update An Entity Token](#update-an-entitys-token)
 
 ## Download Protocol Buffer File
 
@@ -486,7 +489,7 @@ This method retrieves the stored tokens for a given entity.
 
 ```bash
 grpcurl -plaintext \
-    -d '{"long_lived_token": "entity_id:long_lived_token"}' \
+    -d '{"long_lived_token": "long_lived_token"}' \
     -proto protos/v1/vault.proto \
 localhost:6000 vault.v1.Entity/ListEntityStoredTokens
 ```
@@ -603,3 +606,232 @@ localhost:6000 vault.v1.Entity/StoreEntityToken <payload.json
 	"success": true
 }
 ```
+
+#### Get Entity Access Token And Decrypt Payload
+
+This function retrieves an entity's access token and decrypts the payload.
+
+---
+
+> `request` **GetEntityAccessTokenAndDecryptPayloadRequest**
+
+> [!IMPORTANT]
+>
+> The table lists only the required fields for this step. Other fields will be
+> ignored.
+
+| Field              | Type   | Description                                                         |
+| ------------------ | ------ | ------------------------------------------------------------------- |
+| device_id          | string | The unique identifier of the device used by the entity.             |
+| payload_ciphertext | string | The encrypted payload ciphertext that needs to be decrypted.        |
+| platform           | string | The platform from which the token is being issued. (e.g., "gmail"). |
+
+---
+
+> `response` **GetEntityAccessTokenAndDecryptPayloadResponse**
+
+> [!IMPORTANT]
+>
+> The table lists only the fields that are populated for this step. Other fields
+> may be empty, omitted, or false.
+
+| Field             | Type   | Description                                                                |
+| ----------------- | ------ | -------------------------------------------------------------------------- |
+| message           | string | A response message from the server.                                        |
+| success           | bool   | Indicates if the operation was successful.                                 |
+| payload_plaintext | string | The decrypted payload plaintext.                                           |
+| token             | string | The retrieved token associated with the entity for the specified platform. |
+
+---
+
+> `method` **GetEntityAccessTokenAndDecryptPayload**
+
+> [!TIP]
+>
+> The examples below use
+> [grpcurl](https://github.com/fullstorydev/grpcurl#grpcurl).
+
+> [!NOTE]
+>
+> Here is what a successful response from the server looks like.
+>
+> The server would return a status code of `0 OK` if the API transaction goes
+> through without any friction. Otherwise, it will return any other code out of
+> the
+> [17 codes supported by gRPC](https://grpc.github.io/grpc/core/md_doc_statuscodes.html).
+
+---
+
+**Sample request**
+
+```bash
+grpcurl -plaintext \
+    -d '{"device_id": "device_id", "payload_ciphertext": "encrypted_payload", "platform": "gmail"}' \
+    -proto protos/v1/vault.proto \
+localhost:6000 vault.v1.Entity/GetEntityAccessTokenAndDecryptPayload
+```
+
+---
+
+**Sample response**
+
+```json
+{
+	"message": "Successfully fetched tokens and decrypted payload",
+	"success": true,
+	"payload_plaintext": "Decrypted payload content",
+	"token": "retrieved_token"
+}
+```
+
+---
+
+#### Encrypt Payload
+
+This function handles the encryption of payload content.
+
+---
+
+> `request` **EncryptPayloadRequest**
+
+> [!IMPORTANT]
+>
+> The table lists only the required fields for this step. Other fields will be
+> ignored.
+
+| Field             | Type   | Description                                             |
+| ----------------- | ------ | ------------------------------------------------------- |
+| device_id         | string | The unique identifier of the device used by the entity. |
+| payload_plaintext | string | The plaintext payload content to be encrypted.          |
+
+---
+
+> `response` **EncryptPayloadResponse**
+
+> [!IMPORTANT]
+>
+> The table lists only the fields that are populated for this step. Other fields
+> may be empty, omitted, or false.
+
+| Field              | Type   | Description                                |
+| ------------------ | ------ | ------------------------------------------ |
+| message            | string | A response message from the server.        |
+| payload_ciphertext | string | The encrypted payload ciphertext.          |
+| success            | bool   | Indicates if the operation was successful. |
+
+---
+
+> `method` **EncryptPayload**
+
+> [!TIP]
+>
+> The examples below use
+> [grpcurl](https://github.com/fullstorydev/grpcurl#grpcurl).
+
+> [!NOTE]
+>
+> Here is what a successful response from the server looks like.
+>
+> The server would return a status code of `0 OK` if the API transaction goes
+> through without any friction. Otherwise, it will return any other code out of
+> the
+> [17 codes supported by gRPC](https://grpc.github.io/grpc/core/md_doc_statuscodes.html).
+
+---
+
+**Sample request**
+
+```bash
+grpcurl -plaintext \
+    -d '{"device_id": "device_id", "payload_plaintext": "plaintext_payload"}' \
+    -proto protos/v1/vault.proto \
+localhost:6000 vault.v1.Entity/EncryptPayload
+```
+
+---
+
+**Sample response**
+
+```json
+{
+	"message": "Successfully encrypted payload.",
+	"payload_ciphertext": "encrypted_payload",
+	"success": true
+}
+```
+
+---
+
+#### Update An Entity's Token
+
+This function updates tokens associated with an entity.
+
+---
+
+> `request` **UpdateEntityTokenRequest**
+
+> [!IMPORTANT]
+>
+> The table lists only the required fields for this step. Other fields will be
+> ignored.
+
+| Field              | Type   | Description                                                          |
+| ------------------ | ------ | -------------------------------------------------------------------- |
+| device_id          | string | The unique identifier of the device used by the entity.              |
+| token              | string | The new token to be updated for the entity.                          |
+| platform           | string | The platform from which the token is being updated. (e.g., "gmail"). |
+| account_identifier | string | The identifier of the account associated with the token.             |
+
+---
+
+> `response` **UpdateEntityTokenResponse**
+
+> [!IMPORTANT]
+>
+> The table lists only the fields that are populated for this step. Other fields
+> may be empty, omitted, or false.
+
+| Field   | Type   | Description                                |
+| ------- | ------ | ------------------------------------------ |
+| message | string | A response message from the server.        |
+| success | bool   | Indicates if the operation was successful. |
+
+---
+
+> `method` **UpdateEntityToken**
+
+> [!TIP]
+>
+> The examples below use
+> [grpcurl](https://github.com/fullstorydev/grpcurl#grpcurl).
+
+> [!NOTE]
+>
+> Here is what a successful response from the server looks like.
+>
+> The server would return a status code of `0 OK` if the API transaction goes
+> through without any friction. Otherwise, it will return any other code out of
+> the
+> [17 codes supported by gRPC](https://grpc.github.io/grpc/core/md_doc_statuscodes.html).
+
+---
+
+**Sample request**
+
+```bash
+grpcurl -plaintext \
+    -d '{"device_id": "device_id", "token": "new_token", "platform": "gmail", "account_identifier": "account_id"}' \
+    -proto protos/v1/vault.proto \
+localhost:6000 vault.v1.Entity/UpdateEntityToken
+```
+
+---
+
+**Sample response**
+
+```json
+{
+	"message": "Token updated successfully.",
+	"success": true
+}
+```

From 98502d6cd82d0eb0db74870ac99ede644b8fcf91 Mon Sep 17 00:00:00 2001
From: Promise Fru <info@promisefru.com>
Date: Mon, 17 Jun 2024 15:07:59 +0100
Subject: [PATCH 6/6] docs: add Cryptographic Methods Used in the Vault

---
 docs/security.md | 43 +++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 41 insertions(+), 2 deletions(-)

diff --git a/docs/security.md b/docs/security.md
index 3c77c01..3c4350e 100644
--- a/docs/security.md
+++ b/docs/security.md
@@ -2,8 +2,47 @@
 
 ## Password Security
 
-Passwords are secured using HMAC-512. HMAC ([Hash-based Message Authentication Code](https://en.wikipedia.org/wiki/HMAC)) is a MAC defined in [RFC2104](http://www.ietf.org/rfc/rfc2104.txt) and [FIPS-198](http://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.198-1.pdf) and constructed using a cryptographic hash algorithm.
+Passwords are secured using HMAC-512. HMAC
+([Hash-based Message Authentication Code](https://en.wikipedia.org/wiki/HMAC))
+is a MAC defined in [RFC2104](http://www.ietf.org/rfc/rfc2104.txt) and
+[FIPS-198](http://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.198-1.pdf) and
+constructed using a cryptographic hash algorithm.
 
 ## Data Security
 
-Data is secured using AES-CBC. AES ([Advanced Encryption Standard](http://en.wikipedia.org/wiki/Advanced_Encryption_Standard)) is a symmetric block cipher standardized by [NIST](http://csrc.nist.gov/publications/fips/fips197/fips-197.pdf) . It has a fixed data block size of 16 bytes. Its keys can be 128, 192, or 256 bits long.
+Data is secured using AES-CBC. AES
+([Advanced Encryption Standard](http://en.wikipedia.org/wiki/Advanced_Encryption_Standard))
+is a symmetric block cipher standardized by
+[NIST](http://csrc.nist.gov/publications/fips/fips197/fips-197.pdf) . It has a
+fixed data block size of 16 bytes. Its keys can be 128, 192, or 256 bits long.
+
+## Cryptographic Methods Used in the Vault
+
+These cryptographic methods are implemented in the [crypto.py](../src/crypto.py)
+file within the vault.
+
+### 1. AES Encryption
+
+**AES (Advanced Encryption Standard)**:
+
+- **Key Size**: 32 bytes
+- **Mode of Operation**: AES.MODE_EAX
+- **Usage**:
+  - Encrypts and decrypts data at rest in the vault.
+
+### 2. HMAC Generation
+
+**HMAC (Hash-based Message Authentication Code)**:
+
+- **Algorithm**: SHA-512
+- **Key Size**: 32 bytes
+- **Usage**:
+  - Generates and verifies HMACs for unique values in the vault.
+
+### 3. Fernet Encryption
+
+**Fernet encryption**:
+
+- **Key Size**: 32 bytes
+- **Usage**:
+  - Encrypts and decrypts identity tokens used by the vault.