DownUnderCTF 2021 | secuchat

#downunderctf2021

import sqlite3
from Crypto.PublicKey import RSA
from math import gcd
from Crypto.Util.number import isPrime
from Crypto.Cipher import PKCS1_OAEP, AES

def dict_factory(cursor, row):
    d = {}
    for idx, col in enumerate(cursor.description):
        d[col[0]] = row[idx]
    return d

con = sqlite3.connect("secuchat.db")
con.row_factory = dict_factory
cur = con.cursor()

# {"username": "key", ...} dictionary
users = {}

keys = {}
for user in cur.execute("SELECT * FROM User"):
    key = RSA.importKey(user['rsa_key'])
    keys[user["username"]] = {"n": key.n, "p": None, "q": None, "name": user["username"], "e": 65537}

for i, u in keys.items():
    for j, v in keys.items():
        if i == j:
            continue

        g = gcd(keys[i]["n"], keys[j]["n"])
        if g != 1:
            assert isPrime(g), g
            keys[i]["p"] = g
            keys[i]["q"] = keys[i]["n"] // g
            keys[i]["d"] = pow(65537, -1, (keys[i]["p"] - 1)*(keys[i]["q"] - 1))

            keys[j]["p"] = g
            keys[j]["q"] = keys[j]["n"] // g
            keys[j]["d"] = pow(65537, -1, (keys[j]["p"] - 1)*(keys[j]["q"] - 1))

users = []
for k, v in keys.items():
    if v["p"]:
        users.append(k)

src, dst = users

convs = list(cur.execute("SELECT * from Conversation where initiator = ? OR peer = ?", (src, dst)))
for c in convs:
    print(c)
    param_id = c["initial_parameters"]
    params = list(cur.execute("SELECT * from Parameters where id = ?", (param_id,)))[0]

    if c["initiator"] == src:
        k = "encrypted_aes_key_for_initiator"
        x = keys[src]

    if c["initiator"] == dst:
        k = "encrypted_aes_key_for_initiator"
        x = keys[dst]

    if c["peer"] == src:
        k = "encrypted_aes_key_for_peer"
        x = keys[src]

    if c["peer"] == dst:
        k = "encrypted_aes_key_for_peer"
        x = keys[dst]

    oaep = PKCS1_OAEP.new(RSA.construct((x["n"], x["e"], x["d"], x["p"], x["q"])))
    key = oaep.decrypt(params[k])

    msgs = list(cur.execute("SELECT * from Message where conversation = ? ORDER BY timestamp ASC", (c["id"],)))
    for m in msgs:
        aes = AES.new(mode=AES.MODE_CBC, key=key, iv=params["iv"])
        text = aes.decrypt(m["encrypted_message"])
        print(text)

        param_id = m["next_parameters"]
        params = list(cur.execute("SELECT * from Parameters where id = ?", (param_id,)))[0]

        oaep = PKCS1_OAEP.new(RSA.construct((x["n"], x["e"], x["d"], x["p"], x["q"])))
        key = oaep.decrypt(params[k])