yearning-cli 0.1.5__py3-none-any.whl

mysql_proxy/README.md

@@ -0,0 +1,32 @@
+# MySQL 协议代理
+
+MySQL 协议代理模块，让 DBeaver、Navicat、MySQL CLI 等客户端通过 Yearning 审计平台执行查询。
+
+通过 `sql proxy` 子命令启动。详见 [主 README](../README.md#mysql-协议代理)。
+
+## 原理
+
+```
+┌──────────────┐   MySQL协议    ┌──────────────┐   HTTP/WS      ┌──────────────┐
+│  MySQL 客户端 │ ────────────▶ │  sql proxy   │ ────────────▶ │  Yearning    │
+│  (Navicat等) │               │  (3307端口)   │                │  (审计平台)   │
+└──────────────┘               └──────────────┘                └──────────────┘
+```
+
+代理服务器默认监听 `0.0.0.0:3307`，将 MySQL 协议转换为 Yearning API 调用。需要仅允许本机访问时可启动 `sql proxy --host 127.0.0.1`；需要 IPv6 全网卡监听时可使用 `sql proxy --host ::`。
+
+## 数据源切换
+
+连接用户名 = 数据源名称。免密码登录，自动解析数据源：
+
+| 用户名 | 行为 |
+|--------|------|
+| `数据源名称` | 精确匹配该数据源 |
+| `root` / 任意不存在的名称 | 回退到 `~/.sqlrc` 配置的默认数据源 |
+
+密码字段无需填写，任意值均可。
+
+## 前提条件
+
+1. 已运行 `sql login` 登录 Yearning，或已配置 `YEARNING_USER`/`YEARNING_PASS`、`~/.sqlrc` 的 `yearning_user`/`yearning_pass` 供 token 失效时自动登录
+2. `~/.sqlrc` 配置了默认数据源（用于用户名不匹配时的回退）

mysql_proxy/__init__.py

@@ -0,0 +1 @@
+"""MySQL protocol proxy — connect MySQL clients to Yearning audit platform."""

mysql_proxy/protocol.py

@@ -0,0 +1,342 @@
+"""MySQL protocol encoding/decoding utilities."""
+
+import hashlib
+import struct
+
+# MySQL column types
+MYSQL_TYPE_VARCHAR = 0xFC
+MYSQL_TYPE_LONG = 0x03
+MYSQL_TYPE_LONGLONG = 0x08
+MYSQL_TYPE_DOUBLE = 0x05
+MYSQL_TYPE_DECIMAL = 0x00
+MYSQL_TYPE_DATETIME = 0x0C
+MYSQL_TYPE_DATE = 0x0A
+MYSQL_TYPE_TIMESTAMP = 0x07
+MYSQL_TYPE_VAR_STRING = 0xFD
+MYSQL_TYPE_STRING = 0xFE
+MYSQL_TYPE_BLOB = 0xFC
+MYSQL_TYPE_NULL = 0x06
+
+# Capability flags
+CLIENT_LONG_PASSWORD = 0x00000001
+CLIENT_FOUND_ROWS = 0x00000002
+CLIENT_LONG_FLAG = 0x00000004
+CLIENT_CONNECT_WITH_DB = 0x00000008
+CLIENT_PROTOCOL_41 = 0x00000200
+CLIENT_INTERACTIVE = 0x00000400
+CLIENT_TRANSACTIONS = 0x00002000
+CLIENT_SECURE_CONNECTION = 0x00008000
+CLIENT_MULTI_STATEMENTS = 0x00010000
+CLIENT_MULTI_RESULTS = 0x00020000
+CLIENT_PLUGIN_AUTH = 0x00080000
+CLIENT_CONNECT_ATTRS = 0x00100000
+CLIENT_PLUGIN_AUTH_LENENC_CLIENT_DATA = 0x00200000
+
+SERVER_STATUS_AUTOCOMMIT = 0x0002
+SERVER_STATUS_NO_INDEX_USED = 0x0010
+
+DEFAULT_CAPABILITIES = (
+    CLIENT_LONG_PASSWORD
+    | CLIENT_FOUND_ROWS
+    | CLIENT_LONG_FLAG
+    | CLIENT_CONNECT_WITH_DB
+    | CLIENT_PROTOCOL_41
+    | CLIENT_INTERACTIVE
+    | CLIENT_TRANSACTIONS
+    | CLIENT_SECURE_CONNECTION
+    | CLIENT_MULTI_STATEMENTS
+    | CLIENT_MULTI_RESULTS
+    | CLIENT_PLUGIN_AUTH
+    | CLIENT_PLUGIN_AUTH_LENENC_CLIENT_DATA
+)
+
+
+def make_packet(payload: bytes, seq: int) -> bytes:
+    """Wrap payload in MySQL packet format: 3-byte length + 1-byte seq + payload.
+
+    Sequence number wraps at 256 per MySQL protocol spec.
+    """
+    length = len(payload)
+    header = struct.pack("<I", length)[:3] + struct.pack("B", seq % 256)
+    return header + payload
+
+
+def parse_packet(data: bytes):
+    """Parse MySQL packet. Returns (payload, seq_id, remaining_data)."""
+    if len(data) < 4:
+        return None, None, data
+    length = struct.unpack("<I", data[:3] + b"\x00")[0]
+    seq = data[3]
+    if len(data) < 4 + length:
+        return None, None, data
+    payload = data[4 : 4 + length]
+    remaining = data[4 + length :]
+    return payload, seq, remaining
+
+
+def read_null_terminated_string(data: bytes, offset: int = 0):
+    """Read a null-terminated string from bytes."""
+    end = data.index(b"\x00", offset)
+    return data[offset:end].decode("utf-8", errors="replace"), end + 1
+
+
+def read_lenenc_int(data: bytes, offset: int = 0):
+    """Read a length-encoded integer."""
+    first = data[offset]
+    if first < 0xFB:
+        return first, offset + 1
+    elif first == 0xFC:
+        return struct.unpack("<H", data[offset + 1 : offset + 3])[0], offset + 3
+    elif first == 0xFD:
+        return struct.unpack("<I", data[offset + 1 : offset + 4] + b"\x00")[0], offset + 4
+    elif first == 0xFE:
+        return struct.unpack("<Q", data[offset + 1 : offset + 9])[0], offset + 9
+    else:
+        return 0, offset + 1
+
+
+def write_lenenc_int(value: int) -> bytes:
+    """Encode an integer as length-encoded."""
+    if value < 0xFB:
+        return struct.pack("B", value)
+    elif value < 0x10000:
+        return b"\xfc" + struct.pack("<H", value)
+    elif value < 0x1000000:
+        return b"\xfd" + struct.pack("<I", value)[:3]
+    else:
+        return b"\xfe" + struct.pack("<Q", value)
+
+
+def write_lenenc_str(value: bytes) -> bytes:
+    """Encode bytes as length-encoded string."""
+    return write_lenenc_int(len(value)) + value
+
+
+def make_handshake(connection_id: int, salt: bytes, server_version: str = "8.0.0-proxy") -> bytes:
+    """Create MySQL handshake V10 packet."""
+    payload = bytearray()
+    payload.append(0x0A)  # protocol version
+    payload.extend(server_version.encode("utf-8") + b"\x00")
+    payload.extend(struct.pack("<I", connection_id))
+    payload.extend(salt[:8])  # auth_plugin_data_part_1
+    payload.append(0x00)  # filler
+
+    cap_lower = DEFAULT_CAPABILITIES & 0xFFFF
+    cap_upper = (DEFAULT_CAPABILITIES >> 16) & 0xFFFF
+    payload.extend(struct.pack("<H", cap_lower))
+    payload.append(0x21)  # utf8mb4 charset
+    payload.extend(struct.pack("<H", SERVER_STATUS_AUTOCOMMIT))
+    payload.extend(struct.pack("<H", cap_upper))
+    payload.append(21)  # auth_plugin_data_length (8 + 12 + 1 null = 21)
+    payload.extend(b"\x00" * 10)  # reserved
+
+    # auth_plugin_data_part_2 (12 bytes + null terminator)
+    payload.extend(salt[8:20] + b"\x00")
+    payload.extend(b"mysql_native_password\x00")
+
+    return make_packet(bytes(payload), 0)
+
+
+def make_ok(affected_rows=0, last_id=0, status=SERVER_STATUS_AUTOCOMMIT, warnings=0, message="", seq=1):
+    """Create OK packet."""
+    payload = bytearray()
+    payload.append(0x00)
+    payload.extend(write_lenenc_int(affected_rows))
+    payload.extend(write_lenenc_int(last_id))
+    payload.extend(struct.pack("<H", status))
+    payload.extend(struct.pack("<H", warnings))
+    payload.extend(message.encode("utf-8"))
+    return make_packet(bytes(payload), seq)
+
+
+def make_err(error_code=1045, message="Access denied", sql_state="HY000", seq=1):
+    """Create ERR packet."""
+    payload = bytearray()
+    payload.append(0xFF)
+    payload.extend(struct.pack("<H", error_code))
+    payload.append(0x23)  # '#' marker
+    payload.extend(sql_state.encode("utf-8"))
+    payload.extend(message.encode("utf-8"))
+    return make_packet(bytes(payload), seq)
+
+
+def make_eof(warnings=0, status=SERVER_STATUS_AUTOCOMMIT, seq=1):
+    """Create EOF packet."""
+    payload = bytearray()
+    payload.append(0xFE)
+    payload.extend(struct.pack("<H", warnings))
+    payload.extend(struct.pack("<H", status))
+    return make_packet(bytes(payload), seq)
+
+
+def make_column_def(name: str, col_type=MYSQL_TYPE_VARCHAR, seq=0):
+    """Create column definition packet."""
+    payload = bytearray()
+    payload.extend(write_lenenc_str(b"def"))  # catalog
+    payload.extend(write_lenenc_str(b""))  # schema
+    payload.extend(write_lenenc_str(b""))  # table
+    payload.extend(write_lenenc_str(b""))  # org_table
+    payload.extend(write_lenenc_str(name.encode("utf-8")))  # name
+    payload.extend(write_lenenc_str(name.encode("utf-8")))  # org_name
+    payload.append(0x0C)  # next_length (fixed 12)
+    payload.extend(struct.pack("<H", 0x21))  # utf8mb4 charset
+    payload.extend(struct.pack("<I", 256))  # column length
+    payload.append(col_type)
+    payload.extend(struct.pack("<H", 0x0000))  # flags
+    payload.append(0x00)  # decimals
+    payload.extend(b"\x00\x00")  # filler
+    return make_packet(bytes(payload), seq)
+
+
+def make_row(values: list, seq=0):
+    """Create row data packet. Each value is lenenc-string or NULL (0xFB)."""
+    payload = bytearray()
+    for v in values:
+        if v is None:
+            payload.append(0xFB)
+        else:
+            s = str(v).encode("utf-8")
+            payload.extend(write_lenenc_str(s))
+    return make_packet(bytes(payload), seq)
+
+
+def parse_handshake_response(payload: bytes):
+    """Parse client handshake response. Returns (username, auth_response, database)."""
+    offset = 0
+    # capability flags (4 bytes)
+    capabilities = struct.unpack("<I", payload[offset : offset + 4])[0]
+    offset += 4
+    # max packet size (4 bytes)
+    offset += 4
+    # character set (1 byte)
+    offset += 1
+    # reserved (23 bytes)
+    offset += 23
+    # username (null-terminated)
+    username, offset = read_null_terminated_string(payload, offset)
+    # auth response
+    auth_response = b""
+    if capabilities & CLIENT_PLUGIN_AUTH_LENENC_CLIENT_DATA:
+        auth_len, offset = read_lenenc_int(payload, offset)
+        auth_response = payload[offset : offset + auth_len]
+        offset += auth_len
+    elif capabilities & CLIENT_SECURE_CONNECTION:
+        auth_len = payload[offset]
+        offset += 1
+        auth_response = payload[offset : offset + auth_len]
+        offset += auth_len
+    # database
+    database = ""
+    if capabilities & CLIENT_CONNECT_WITH_DB and offset < len(payload):
+        try:
+            database, offset = read_null_terminated_string(payload, offset)
+        except ValueError:
+            pass
+    return username, auth_response, database
+
+
+def verify_password(salt: bytes, auth_response: bytes, password: str) -> bool:
+    """Verify mysql_native_password authentication."""
+    if not password:
+        return auth_response == b""
+    stage1 = hashlib.sha1(password.encode("utf-8")).digest()
+    stage2 = hashlib.sha1(stage1).digest()
+    combined = hashlib.sha1(salt + stage2).digest()
+    expected = bytes(a ^ b for a, b in zip(combined, stage1))
+    return expected == auth_response
+
+
+def infer_column_type(value) -> int:
+    """Infer MySQL column type from Python value."""
+    if value is None:
+        return MYSQL_TYPE_NULL
+    if isinstance(value, bool):
+        return MYSQL_TYPE_LONG
+    if isinstance(value, int):
+        return MYSQL_TYPE_LONGLONG
+    if isinstance(value, float):
+        return MYSQL_TYPE_DOUBLE
+    return MYSQL_TYPE_VARCHAR
+
+
+def make_result_set(columns: list, rows: list, seq_start=1):
+    """Build complete result set packets. Returns list of (packet_bytes, seq).
+
+    Kept for backward compatibility (small result sets, simple results).
+    For large result sets, use stream_result_set instead.
+    """
+    packets = []
+    seq = seq_start
+
+    # Column count
+    packets.append((make_packet(write_lenenc_int(len(columns)), seq), seq))
+    seq += 1
+
+    # Column definitions
+    for i, col in enumerate(columns):
+        col_name = str(col)
+        # Infer type from first row
+        col_type = MYSQL_TYPE_VARCHAR
+        if rows:
+            val = rows[0][i] if i < len(rows[0]) else None
+            col_type = infer_column_type(val)
+        packets.append((make_column_def(col_name, col_type, seq), seq))
+        seq += 1
+
+    # EOF after columns
+    packets.append((make_eof(seq=seq), seq))
+    seq += 1
+
+    # Row data
+    for row in rows:
+        values = [row[i] if i < len(row) else None for i in range(len(columns))]
+        packets.append((make_row(values, seq), seq))
+        seq += 1
+
+    # EOF after rows
+    status = SERVER_STATUS_AUTOCOMMIT | SERVER_STATUS_NO_INDEX_USED
+    packets.append((make_eof(status=status, seq=seq), seq))
+    seq += 1
+
+    return packets
+
+
+def make_result_header(columns: list, rows_sample: list, seq_start=1):
+    """Build column count + column definitions + EOF header packets.
+
+    Returns (list_of_packets, next_seq). Use with stream_rows for large result sets.
+    """
+    packets = []
+    seq = seq_start
+
+    packets.append((make_packet(write_lenenc_int(len(columns)), seq), seq))
+    seq += 1
+
+    for i, col in enumerate(columns):
+        col_name = str(col)
+        col_type = MYSQL_TYPE_VARCHAR
+        if rows_sample:
+            val = rows_sample[0][i] if i < len(rows_sample[0]) else None
+            col_type = infer_column_type(val)
+        packets.append((make_column_def(col_name, col_type, seq), seq))
+        seq += 1
+
+    packets.append((make_eof(seq=seq), seq))
+    seq += 1
+
+    return packets, seq
+
+
+def stream_rows(columns: list, rows: list, seq_start: int):
+    """Yield row packets one by one, then final EOF packet.
+
+    This avoids building all row packets in memory at once.
+    """
+    seq = seq_start
+    for row in rows:
+        values = [row[i] if i < len(row) else None for i in range(len(columns))]
+        yield make_row(values, seq), seq
+        seq += 1
+
+    status = SERVER_STATUS_AUTOCOMMIT | SERVER_STATUS_NO_INDEX_USED
+    yield make_eof(status=status, seq=seq), seq