#!/usr/bin/env python3

# This code can be used to solve the crypto 01 challenge for the CTF
# organized by RomHack and HTB (2021)
#
# Author: Leonardo Tamiano
# Date: 2021-09-30
#

import socket
import binascii
from itertools import cycle
from Crypto.Util.number import bytes_to_long, inverse
from Crypto.Hash import SHA

HOST = '127.0.0.1'
PORT = 23333
KEY_LENGTH = 12

# -------------------------

def extract_public_params(sock):
    sock.send(b'3')

    data = sock.recv(6000).split(b'\n')
    p = data[1].split(b"=")[1].strip().decode("ascii")
    q = data[2].split(b"=")[1].strip().decode("ascii")
    g = data[3].split(b"=")[1].strip().decode("ascii")
    y = data[4].split(b"=")[1].strip().decode("ascii")

    return int(p, 0), int(q, 0), int(g, 0), int(y, 0)

# -------------------------

def extract_key(sock):
    global KEY_LENGTH
    
    plaintext_bytes = b"A" * KEY_LENGTH

    # -- ask the server to sign plaintext
    sock.send(b'1')
    sock.recv(1024)
    sock.send(plaintext_bytes)

    # -- get back the ciphered password
    data = sock.recv(1024)
    ciphertext = data.split(b"\n")[1]
    
    # -- do a KNOWN-PLAINTEXT-ATTACK and get back the key
    ciphertext_bytes = binascii.unhexlify(ciphertext)
    key_bytes = [int(a ^ b) for (a, b) in zip(plaintext_bytes, ciphertext_bytes)]
    return key_bytes

# -------------------------

def sign(sock, msg):
    
    # -- ask for sign
    sock.send(b'1')
    sock.recv(1024)
    sock.send(msg)

    # -- get sign
    
    data = sock.recv(1024)
    data = data.split(b"\n")[3].split(b",")

    r = "0x" + data[0].decode("ascii")
    s = "0x" + data[1].decode("ascii")

    return r, s

# -------------------------

def verify(sock, msg, r, s):
    # -- ask for verify
    sock.send(b'2')
    sock.recv(1024)

    sock.send(msg)
    sock.recv(1024)

    sock.send(r)
    sock.recv(1024)

    sock.send(s)

    return sock.recv(1024)
    
# -------------------------

def exploit():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.connect((HOST, PORT))
        
        # -- ignore initial response
        sock.recv(1024)

        # -- 0) get server public parameters
        p, q, g, y = extract_public_params(sock)

        print("\n[INFO]: Extract params from server:\n"
              f"\tp = {p}\n"
              f"\tq = {q}\n"
              f"\tg = {g}\n"
              f"\ty = {y}\n")
        
        # -- 1) extract key with XOR
        key = extract_key(sock)

        print("[INFO]: Extracted key from server:\n"
              f"\tkey = {key}\n")
        
        # -- 2) sign a message and get (r, s)
        msg = b"HELLO"
        r, s = sign(sock, msg)
        r, s = int(r, 0), int(s, 0)

        print(f"[INFO]: Signed message m = {msg}:\n"
              f"\tr = {r}\n"
              f"\ts = {s}\n")
        
        # -- 3) extract x
        h = bytes_to_long(SHA.new(msg).digest())
        k = bytes_to_long(bytes([a ^ b for (a, b) in zip(msg, cycle(key))])) % q
        x = inverse(r, q) * (s * k - h) % q

        print("[INFO]: Extracted secret key from server:\n"
              f"\tx = {x}\n")

        # -- 4) sign message "give me flag"
        msg = b"give me flag"
        h = bytes_to_long(SHA.new(msg).digest())
        k = bytes_to_long(bytes([a ^ b for (a, b) in zip(msg, cycle(key))])) % q

        r = pow(g, k, p) % q
        s = (inverse(k, q) * (h + x * r)) % q

        print(f"[INFO]: Signed message m = {msg}:\n"
              f"\tr = {hex(r)}\n"
              f"\ts = {hex(s)}\n")

        # -- 5) verify the sign and get the flag
        flag = verify(sock, msg, hex(r)[2:].encode(), hex(s)[2:].encode())

        print("[INFO]: Got the flag:\n"
              f"\tflag = {flag}\n")



if __name__ == "__main__":
    exploit()