# sageの環境構築がめんどうなので，https://sagecell.sagemath.org/ で実行する
# poly1305のtag復元は https://github.com/tl2cents/AEAD-Nonce-Reuse-Attacks/blob/main/chacha-poly1305/chacha_poly1305_forgery.py を利用した．
# ChaCha20での暗号化の部分も上記OSSで実装されているが，本コードは自分の実装をmain関数に記述した

from sage.all import GF, ZZ, PolynomialRing

# the implementation of the Poly1305 in the standard lib is as follows:
def poly1305(msg:bytes, key:bytes, byteorder='little'):
    """ A pure python implementation of the Poly1305 MAC function
    Reference: https://datatracker.ietf.org/doc/html/rfc7539#section-2.5.1
    Args:
        msg (bytes): The message to authenticate
        key (bytes): The 32 byte key to use
    Returns:
        bytes: The 16 byte MAC
    """
    p = 2**130 - 5 # the prime number used in Poly1305
    r = int.from_bytes(key[:16], byteorder)
    r = r & 0x0ffffffc0ffffffc0ffffffc0fffffff
    s = int.from_bytes(key[16:], byteorder)
    acc = 0
    for i in range(0, len(msg), 16):
        block = msg[i:i+16] + b'\x01'
        block = int.from_bytes(block, byteorder)
        acc = (acc + block) * r % p
    acc = (acc + s) # no more % p here !!!
    acc = int(acc % 2**128)
    return acc.to_bytes(16, byteorder)

def construct_poly1305_coeffs(msg:bytes, byteorder='little'):
    # get coefficients of the polynomial from the message bytes
    coeff = []
    for i in range(0, len(msg), 16):
        block = msg[i:i+16] + b'\x01'
        block = int.from_bytes(block, byteorder)
        coeff.append(block)
    return coeff

def is_valid_r(r):
    # check if r is a valid Poly1305 key
    # from RFC specification: https://datatracker.ietf.org/doc/html/rfc7539#section-2.5
    return (r & 0x0ffffffc0ffffffc0ffffffc0fffffff) == r

def recover_poly1305_key_from_nonce_reuse(msg1:bytes, tag1:bytes, 
                                          msg2:bytes, tag2:bytes,
                                          byteorder='little'):
    """Recover the Poly1305 key when nonce is reused.
    
    Args:
        msg1 (bytes): The first message
        tag1 (bytes): The first tag
        msg2 (bytes): The second message
        tag2 (bytes): The second tag
    Returns:
        list: keys r, s
    """
    p = 2**130 - 5 # the prime number used in Poly1305
    pp = 2**128
    PolynomialRing_GF = PolynomialRing(GF(p), 'x')
    x = PolynomialRing_GF.gen()
    poly1 = x * PolynomialRing_GF(construct_poly1305_coeffs(msg1, byteorder)[::-1])
    a1 = int.from_bytes(tag1, byteorder)
    poly2 = x * PolynomialRing_GF(construct_poly1305_coeffs(msg2, byteorder)[::-1])
    a2 = int.from_bytes(tag2, byteorder)
    # find all possible keys
    roots = []
    print(f"[+] Searching for the key with poly.degree() = {(poly1 - poly2).degree()}")
    # the range is larger than expected
    # since in poly1305 acc + s does not have to be modulo p
    for tag1 in range(a1, p + pp, pp):
        for tag2 in range(a2, p + pp, pp):
            f = (poly1 - poly2) - (tag1 - tag2)
            root = f.roots(multiplicities=False)
            for r in root:
                r = int(r)
                if is_valid_r(r):
                    # print(f"[+] Found a valid key: {r}")
                    # check if the key is valid
                    s = int(a1 - int(poly1(r))) % pp
                    if (int(poly1(r)) + s) % pp == a1 and (int(poly2(r)) + s) % pp == a2:
                        roots.append((r, s))
    return list(set(roots))

    
def chachapoly1305_merger(ad:bytes, ct:bytes):
    """Merge the associated data and the ciphertext
    
    Args:
        ad (bytes): The associated data
        ct (bytes): The ciphertext
    Returns:
        bytes: The merged data
    """
    def pad(data, block_size=16):
        if len(data) % block_size == 0:
            return data
        return data + b'\x00' * (block_size - len(data) % block_size)
    la = len(ad)
    lc = len(ct)
    out = pad(ad) + pad(ct) + la.to_bytes(8, 'little') + lc.to_bytes(8, 'little')
    return out

def chachapoly1305_nonce_reuse_attack(ad1:bytes, ct1:bytes, tag1:bytes, 
                                      ad2:bytes, ct2:bytes, tag2:bytes):
    """Recover the Chacha-Poly1305 key when nonce is reused.
    
    The two messages are encrypted with the same nonce and the same key.
    Args:
        ad1 (bytes): The first associated data
        ct1 (bytes): The first ciphertext
        tag1 (bytes): The first tag
        ad2 (bytes): The second associated data
        ct2 (bytes): The second ciphertext
        tag2 (bytes): The second tag
    Returns:
        list: keys r, s
    """
    inp1 = chachapoly1305_merger(ad1, ct1)
    inp2 = chachapoly1305_merger(ad2, ct2)
    return recover_poly1305_key_from_nonce_reuse(inp1, tag1, inp2, tag2)

def forge_poly1305_tag(ad:bytes, ct:bytes, r:int, s:int):
    """Forge a Poly1305 tag given the message and the key
    
    Args:
        ad (bytes): The associated data
        ct (bytes): The ciphertext
        r (int): The r key
        s (int): The s key
    """
    key = int(r).to_bytes(16, 'little') + int(s).to_bytes(16, 'little')
    msg = chachapoly1305_merger(ad, ct)
    return poly1305(msg, key)

def chachapoly1305_forgery_attack(ad1:bytes, ct1:bytes, tag1:bytes, 
                                  ad2:bytes, ct2:bytes, tag2:bytes,
                                  target_ct:bytes, target_ad:bytes):
    keys = chachapoly1305_nonce_reuse_attack(ad1, ct1, tag1, ad2, ct2, tag2)
    if len(keys) == 0:
        assert False, "Failed to recover the key for poly1305, probably the nonce is not reused"
    
    for r, s in keys:
        target_tag = forge_poly1305_tag(target_ad, target_ct, r, s)
        return target_tag
    
from binascii import unhexlify, hexlify

def main():
    chiper1 = unhexlify("0b9c1f887aeca7ea98da43a5f4487999304be1fe9149492ecc620df376e75067c19304452bee1535a7f832b8617f0394a2fd625fbbd7329d42cc9cc865c31a2df17de5ec58a7497cc0a091d02d6aafad8e2c83661527bb13253a4b9ca731b7d346bd0d7dd8e4bebecf")
    plain1 = b"Did you know that ChaCha20-Poly1305 is an authenticated encryption algorithm?"

    chiper2 = unhexlify("1b9d1adc23eeb7ab9dc70cbba01c618a2b1fc7f58479012d912648836de34c7691cc5f032bf9503abdb132a17c631fdab7fa651ea6dc22d840d096ce65930122be77a4f955e155b7a7bcc0f2e217396a226eea5e5ed631b7d346bd0d7dd8e4bebecf")

    # keystream を求める（C1 と P1 を使って XOR）(sageでXORは^^でできる．^だとべき乗になってしまうため．)
    keystream = bytes([c1 ^^ p1 for c1, p1 in zip(chiper1, plain1)])

    # goal メッセージ（暗号化したい新しいメッセージ）
    goal = b"But it's only secure if used correctly!"

    # 新しい暗号文を生成（keystream を使って goal を暗号化）
    goal_ciphertext = bytes([g ^^ k for g, k in zip(goal, keystream)])

    # nonce は再利用
    nonce = chiper1[-12:]

    # ct1,2 にはtag+nonceを除く暗号文を入力する
    ad1 = b""  # 今回は利用しない
    ct1 = chiper1[:-28]
    tag1 = chiper1[-28:-12]
    
    ad2 = b""  # 今回は利用しない
    ct2 = chiper2[:-28]
    tag2 = chiper2[-28:-12]

    target_ad = b""  # 今回は利用しない

    # Poly1305タグ復元関数
    target_tag = chachapoly1305_forgery_attack(
        ad1, ct1, tag1, 
        ad2, ct2, tag2, 
        goal_ciphertext,
        target_ad
    )
    
    print("chiper text: ",goal_ciphertext.hex())
    print("Forged tag:", target_tag.hex())
    print("Nonce:", nonce.hex())
    print("")
    print("result: ", goal_ciphertext.hex()+ target_tag.hex()+ nonce.hex())

if __name__=="__main__":
    main()