{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":4117,"databundleVersionId":46665,"sourceType":"competition"},{"sourceId":290230831,"sourceType":"kernelVersion"},{"sourceId":290374285,"sourceType":"kernelVersion"}],"dockerImageVersionId":31236,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# =============================================================================\n# 1. SETUP & INSTALLATION\n# =============================================================================\n!pip install py7zr -q\n\nimport os\nimport glob\nimport shutil\nimport time\nimport re\nimport numpy as np\nimport pandas as pd\nimport cv2\nimport tensorflow as tf\nimport xgboost as xgb\nimport py7zr\n\n# =============================================================================\n# 2. CẤU HÌNH HỆ THỐNG\n# =============================================================================\nINPUT_DIR = \"/kaggle/input/malware-classification\"\nWORK_DIR = \"/kaggle/working/samples_test\"\nLABEL_FILE_SRC = f\"{INPUT_DIR}/trainLabels.csv\"\nLABEL_FILE_DST = \"/kaggle/working/trainLabels.csv\"\n\n# Đường dẫn Model\nCNN_MODEL_PATH = '/kaggle/input/cnn-t4-x2/best_cnn_sync.keras'\nXGB_MODEL_PATH = '/kaggle/input/mao-oe/xgb_out_C/xgb_C_opcode_segment_2nd.json'\n\nCLASSES = [\"Ramnit\", \"Lollipop\", \"Kelihos_ver3\", \"Vundo\", \"Simda\",\n           \"Tracur\", \"Kelihos_ver1\", \"Obfuscator.ACY\", \"Gatak\"]\nIMG_SIZE = (226, 226)\n\n# Danh sách Opcode & Segment phổ biến để trích xuất\nOPCODES = ['mov', 'push', 'pop', 'call', 'add', 'sub', 'inc', 'dec', 'or', 'xor', 'and', 'lea', 'test', 'jmp']\nSECTIONS = ['.text', '.data', '.rdata', '.idata', '.reloc', '.bss', '.code']\n\n# =============================================================================\n# 3. HELPER FUNCTIONS\n# =============================================================================\n\ndef smart_extract(archive_path, dest_dir, num_samples=10):\n    if not archive_path:\n        print(\"❌ Lỗi: Không tìm thấy file train.7z!\")\n        return\n    \n    if os.path.exists(dest_dir) and len(glob.glob(f\"{dest_dir}/*.bytes\")) >= 5:\n        print(f\"✅ Dữ liệu đã có sẵn tại {dest_dir}. Bỏ qua giải nén.\")\n        return\n\n    print(f\"📦 Đang giải nén {num_samples} mẫu từ: {os.path.basename(archive_path)}...\")\n    try:\n        with py7zr.SevenZipFile(archive_path, mode='r') as z:\n            all_files = z.getnames()\n            bytes_files = [f for f in all_files if f.endswith('.bytes')][:num_samples]\n            # Lấy luôn file .asm tương ứng\n            target_files = bytes_files + [f.replace('.bytes', '.asm') for f in bytes_files]\n            \n            z.extract(targets=target_files, path=dest_dir)\n            print(\"✅ Giải nén hoàn tất!\")\n    except Exception as e:\n        print(f\"❌ Lỗi giải nén: {e}\")\n\ndef load_labels(csv_path):\n    if not os.path.exists(csv_path): return {}\n    df = pd.read_csv(csv_path)\n    label_map = {}\n    for _, row in df.iterrows():\n        label_map[row['Id']] = CLASSES[int(row['Class']) - 1]\n    return label_map\n\ndef process_image(path):\n    \"\"\"Xử lý .bytes -> Ảnh cho CNN\"\"\"\n    try:\n        with open(path, 'r') as f:\n            arr = []\n            for line in f:\n                parts = line.strip().split()\n                if len(parts) > 1:\n                    row_bytes = [int(i, 16) if i != '??' else 0 for i in parts[1:]]\n                    arr.append(row_bytes)\n        if not arr: return None\n        img = np.array(arr)\n        side = int(np.sqrt(img.size))\n        if side < 1: side = 1\n        img_2d = img.flatten()[:side*side].reshape(side, side)\n        img_final = cv2.resize(np.uint8(img_2d), IMG_SIZE)\n        return img_final.reshape(1, 226, 226, 1) / 255.0\n    except: return None\n\ndef extract_asm_features(path, expected_dim):\n    \"\"\"\n    Đọc file .asm và đếm Opcode/Segment.\n    Input: Đường dẫn file .asm, số chiều vector model yêu cầu (expected_dim)\n    Output: Vector numpy (1, expected_dim)\n    \"\"\"\n    if not os.path.exists(path):\n        return None\n\n    # Khởi tạo bộ đếm\n    feat_counts = {k: 0 for k in OPCODES + SECTIONS}\n    feat_counts['filesize'] = os.path.getsize(path)\n    \n    try:\n        # Đọc file ASM (Dùng errors='ignore' để tránh lỗi ký tự lạ)\n        with open(path, 'r', errors='ignore') as f:\n            for line in f:\n                line = line.lower().strip()\n                # Đếm Segment\n                for sec in SECTIONS:\n                    if line.startswith(sec):\n                        feat_counts[sec] += 1\n                \n                # Đếm Opcode (Quét đơn giản)\n                # (Trong thực tế cần tách từ kỹ hơn)\n                for op in OPCODES:\n                    if f\" {op} \" in line or f\"\\t{op} \" in line:\n                        feat_counts[op] += 1\n                        \n        # --- CHUYỂN ĐỔI SANG VECTOR ---\n        # Vì ta thiếu vectorizer gốc, ta sẽ tạo một vector có kích thước đúng như model yêu cầu\n        # và điền các giá trị đếm được vào các vị trí đầu tiên (Best Effort).\n        \n        vector = np.zeros((1, expected_dim))\n        \n        # Lấy danh sách giá trị đếm được\n        vals = list(feat_counts.values())\n        \n        # Điền vào vector (Cắt bớt nếu thừa, hoặc để zero nếu thiếu)\n        fill_len = min(len(vals), expected_dim)\n        vector[0, :fill_len] = vals[:fill_len]\n        \n        return vector\n\n    except Exception as e:\n        print(f\"Lỗi đọc ASM: {e}\")\n        return None\n\n# =============================================================================\n# 4. MAIN EXECUTION\n# =============================================================================\ndef main():\n    print(\"🚀 BẮT ĐẦU TEST (CHẾ ĐỘ XỬ LÝ FULL .ASM & .BYTES)\\n\")\n    \n    # 1. Chuẩn bị dữ liệu\n    train_archive = glob.glob(f\"{INPUT_DIR}/**/train.7z\", recursive=True)\n    train_archive = train_archive[0] if train_archive else None\n    smart_extract(train_archive, WORK_DIR, num_samples=15) # Lấy 15 mẫu\n    \n    csv_src = glob.glob(LABEL_FILE_SRC, recursive=True)\n    if csv_src: shutil.copy(csv_src[0], LABEL_FILE_DST)\n\n    # 2. Load Models\n    print(\"\\n📦 Loading Models...\")\n    cnn_model = tf.keras.models.load_model(CNN_MODEL_PATH) if os.path.exists(CNN_MODEL_PATH) else None\n    \n    xgb_model = None\n    if os.path.exists(XGB_MODEL_PATH):\n        xgb_model = xgb.Booster()\n        xgb_model.load_model(XGB_MODEL_PATH)\n        print(f\"✅ XGBoost Loaded (Expects {xgb_model.num_features()} features)\")\n\n    # 3. Inference Loop\n    ground_truth = load_labels(LABEL_FILE_DST)\n    test_files = glob.glob(f\"{WORK_DIR}/**/*.bytes\", recursive=True)\n    \n    if not test_files:\n        print(\"❌ Không tìm thấy file để test.\")\n        return\n\n    print(f\"\\n🔍 Đang phân tích {len(test_files)} file...\\n\")\n    \n    header = f\"{'FILENAME':<25} | {'PREDICTION':<15} | {'TRUE LABEL':<15} | {'CONF':<7} | {'NOTE':<10} | {'TIME'}\"\n    print(\"-\" * 100)\n    print(header)\n    print(\"-\" * 100)\n\n    correct = 0\n    total = 0\n\n    for fpath in test_files:\n        t0 = time.time()\n        fname = os.path.basename(fpath)\n        fid = fname.split('.')[0]\n        asm_path = fpath.replace(\".bytes\", \".asm\")\n        \n        # --- ENSEMBLE ---\n        p_final = np.zeros((1, 9))\n        note = \"\"\n\n        # 1. CNN\n        img = process_image(fpath)\n        if img is not None and cnn_model:\n            p_final += 0.5 * cnn_model.predict(img, verbose=0)\n        \n        # 2. XGBoost (Đã xử lý file ASM)\n        if xgb_model and os.path.exists(asm_path):\n            # Trích xuất đặc trưng từ file ASM thật\n            feats = extract_asm_features(asm_path, xgb_model.num_features())\n            if feats is not None:\n                dmatrix = xgb.DMatrix(feats)\n                p_final += 0.5 * xgb_model.predict(dmatrix).reshape(1, -1)\n                note = \"Full-Ens\"\n            else:\n                note = \"CNN-Only\" # Lỗi đọc asm\n        else:\n             note = \"CNN-Only\" # Không có file asm\n\n        # Result\n        pred = CLASSES[np.argmax(p_final)]\n        conf = np.max(p_final) * 100\n        true_lbl = ground_truth.get(fid, \"Unknown\")\n        \n        # Check correctness\n        status = \"MISS\"\n        if true_lbl != \"Unknown\":\n            total += 1\n            if pred == true_lbl: \n                status = \"OK\"\n                correct += 1\n        else: status = \"?\"\n\n        # Print\n        color = \"\\033[92m\" if status == \"OK\" else (\"\\033[91m\" if status == \"MISS\" else \"\\033[93m\")\n        print(f\"{color}{fname:<25} | {pred:<15} | {true_lbl:<15} | {conf:5.1f}%  | {note:<10} | {time.time()-t0:.2f}s\\033[0m\")\n\n    if total > 0:\n        print(\"-\" * 100)\n        print(f\"✅ Accuracy: {correct/total*100:.2f}% ({correct}/{total})\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-07T17:50:41.664822Z","iopub.execute_input":"2026-01-07T17:50:41.665212Z","iopub.status.idle":"2026-01-07T17:51:28.881204Z","shell.execute_reply.started":"2026-01-07T17:50:41.665180Z","shell.execute_reply":"2026-01-07T17:51:28.880417Z"}},"outputs":[],"execution_count":null}]}