{"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":"none","dataSources":[{"sourceType":"competition","sourceId":4117,"databundleVersionId":46665}],"dockerImageVersionId":31328,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# MS Malware classification: Getting features from raw data\n\n**Author:** Nguyen Le Minh Quan, Phan Ngoc Thuc, Tong Phuc Thien\n\n**Objective:** Extract visual embeddings (Images) and tabular statistics (Byte/ASM frequencies/Opcodes) from ~0.5 TB of raw `.7z` malware files.\n\n**Performance Note:** Optimized for CPU-only execution (Accelerator: None). Leverages RAM disk (`/dev/shm`), Numpy vectorization, and large batch sizes to overcome I/O bottlenecks.","metadata":{}},{"cell_type":"markdown","source":"## Set up and EDA\n\nFirst of all, we should take a general view of upcoming samples like structure of the input, density of 9 families malware","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport os\nimport warnings\n\n# Suppress minor visualization warnings for a clean report\nwarnings.filterwarnings(\"ignore\", category=FutureWarning) \n\n# --- 1. SETUP & INITIALIZATION ---\nbase_path = \"/kaggle/input/competitions/malware-classification/\"\ntrain_labels = pd.read_csv(os.path.join(base_path, \"trainLabels.csv\"))\n\nprint(\"-\" * 50)\nprint(f\"Dataset Overview:\")\nprint(f\"Total training samples: {len(train_labels)}\")\nprint(f\"Total features (Current): {train_labels.shape[1]} (Id, Class)\")\nprint(\"-\" * 50)\n\n# --- 2. CLASS DISTRIBUTION ANALYSIS ---\nclass_counts = train_labels['Class'].value_counts().sort_index()\nclass_names = ['Ramnit', 'Lollipop', 'Kelihos_v3', 'Vundo', 'Simda', \n               'Tracur', 'Kelihos_v1', 'Obfuscator.ACY', 'Gatak']\n\n# Mathematical Quantification of Imbalance\nmax_class = class_counts.max()\nmin_class = class_counts.min()\nimbalance_ratio = max_class / min_class\n\nprint(f\"\\nClass Imbalance Metrics:\")\nprint(f\"Majority Class: {max_class} samples\")\nprint(f\"Minority Class (Simda): {min_class} samples\")\nprint(f\"Imbalance Ratio: {imbalance_ratio:.2f} : 1\")\nprint(\"Conclusion: Severe class imbalance detected. Cost-sensitive learning or weighted sampling is mandatory.\\n\")\n\n# --- 3. VISUALIZATION ---\nplt.figure(figsize=(14, 7))\nsns.set_style(\"whitegrid\")\ncolors = sns.color_palette(\"husl\", 9)\n\n# Added 'hue' and 'legend=False' to comply with modern Seaborn standards\nax = sns.barplot(x=class_counts.index, y=class_counts.values, \n                 palette=colors, hue=class_counts.index, legend=False)\n\n# Add precise value annotations on top of each bar\nfor i, v in enumerate(class_counts.values):\n    ax.text(i, v + 50, str(v), ha='center', fontweight='bold', fontsize=11)\n\nplt.title(\"Malware Family Distribution in Training Set\", fontsize=16, fontweight='bold', pad=20)\nplt.xlabel(\"Malware Family\", fontsize=14)\nplt.ylabel(\"Number of Samples\", fontsize=14)\nplt.xticks(ticks=range(9), labels=class_names, rotation=45, ha='right', fontsize=12)\n\n# Despine to make the chart look cleaner (Tufte's data-ink ratio)\nsns.despine(left=True, bottom=True)\nplt.tight_layout()\nplt.show()\n\n# --- 4. ARCHIVE METADATA ---\ntrain_size_gb = os.path.getsize(os.path.join(base_path, \"train.7z\")) / (1024**3)\ntest_size_gb = os.path.getsize(os.path.join(base_path, \"test.7z\")) / (1024**3)\nprint(\"-\" * 50)\nprint(f\"Storage Footprint:\")\nprint(f\"train.7z Size: {train_size_gb:.2f} GB\")\nprint(f\"test.7z Size:  {test_size_gb:.2f} GB\")\nprint(f\"Total compressed payload: > {train_size_gb + test_size_gb:.2f} GB\")\nprint(\"-\" * 50)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-04-14T17:56:43.250351Z","iopub.execute_input":"2026-04-14T17:56:43.251567Z","iopub.status.idle":"2026-04-14T17:56:43.596811Z","shell.execute_reply.started":"2026-04-14T17:56:43.251519Z","shell.execute_reply":"2026-04-14T17:56:43.595977Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 1. Architecture & Environment Setup\nTo bypass Kaggle's hardware limitations, this pipeline employs a **RAM Disk (`/dev/shm`) strategy** coupled with aggressive `numpy` vectorization.\n\n**Performance Upgrades:**\n* `BATCH_SIZE = 150`: Drastically reduces the number of times `7z` has to scan the 18GB archive header.\n* Added Opcode tracking arrays to extract structural execution logic.\n* Removed hardcoded image resizing to preserve spatial features for Adaptive Pooling later.","metadata":{}},{"cell_type":"code","source":"import os\nimport subprocess\nimport pandas as pd\nimport numpy as np\nfrom collections import Counter\nfrom concurrent.futures import ProcessPoolExecutor, as_completed\nfrom PIL import Image\n\n# --- 1. PIPELINE CONFIGURATION ---\nMODE = 'train' \nARCHIVE_PATH = f\"/kaggle/input/competitions/malware-classification/{MODE}.7z\"\n\nif MODE == 'train':\n    LABEL_PATH = \"/kaggle/input/competitions/malware-classification/trainLabels.csv\"\n    df_targets = pd.read_csv(LABEL_PATH)\n    target_ids = df_targets['Id'].tolist()\nelse:\n    SUBMISSION_PATH = \"/kaggle/input/competitions/malware-classification/sampleSubmission.csv\"\n    df_targets = pd.read_csv(SUBMISSION_PATH)\n    target_ids = df_targets['Id'].tolist()\n\nWORKING_DIR = \"/kaggle/working/\"\nBATCH_DIR = \"/dev/shm/malware_batch\"  \nIMAGE_DIR = f\"/kaggle/working/{MODE}_malware_images\"\n\n# Tối ưu hóa I/O và CPU cho Kaggle Kernel\nBATCH_SIZE = 150      \nMAX_WORKERS = os.cpu_count() or 4      \n\nos.makedirs(IMAGE_DIR, exist_ok=True)\nHEX_CODES = [f\"{i:02x}\".upper() for i in range(256)]\nSEGMENTS = ['.text', '.data', '.bss', '.rdata', '.edata', '.idata', '.rsrc']\n# Core ASM opcodes for execution profiling\nOPCODES = ['mov', 'push', 'pop', 'jmp', 'call', 'ret', 'cmp', 'test', 'add', 'sub', 'inc', 'dec', 'xor', 'jz', 'jnz']","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 2. The Extraction Engine (Worker Node)\nThis section defines the core parsing logic for individual malware files. It is optimized for extreme speed and memory efficiency:\n* **Numpy Vectorization:** Bypasses slow Python list comprehensions using `np.fromiter` and `np.bincount` to parse and aggregate 150MB `.bytes` files instantly.\n* **Lossless Image Rendering:** Converts hexadecimal instructions into grayscale images *without resizing*, keeping the spatial locality completely intact for Global Average Pooling (GAP).\n* **Deep ASM Mining:** Extracts not just segment counts, but also core execution opcodes and obfuscation signatures (`??`).\n* **Immediate Purge Protocol:** Files are permanently deleted from the RAM Disk the exact millisecond their features are extracted.","metadata":{}},{"cell_type":"code","source":"# --- 2. FEATURE EXTRACTION LOGIC ---\ndef get_image_width(size_kb):\n    if size_kb < 10: return 64\n    if size_kb < 30: return 128\n    if size_kb < 60: return 256\n    if size_kb < 100: return 384\n    if size_kb < 200: return 512\n    if size_kb < 500: return 768\n    if size_kb < 1000: return 1024\n    return 2048\n\ndef process_single_file(fid):\n    features = {'Id': fid}\n    byte_path = os.path.join(BATCH_DIR, f\"{fid}.bytes\")\n    asm_path = os.path.join(BATCH_DIR, f\"{fid}.asm\")\n    \n    try:\n        if os.path.exists(byte_path):\n            file_size_kb = os.path.getsize(byte_path) / 1024\n            width = get_image_width(file_size_kb)\n            \n            with open(byte_path, 'r') as f:\n                content = f.read()\n                \n                # EXTRACT OBFUSCATION SIGNATURE\n                features['byte_unknown'] = content.count('??')\n                \n                # Clean content for structural image rendering\n                tokens = content.replace('??', '00').split()\n            \n            # Ép kiểu thần tốc bằng Numpy (Nhanh hơn 30% so với list comprehension)\n            hex_iter = (int(t, 16) for t in tokens if len(t) == 2)\n            int_data = np.fromiter(hex_iter, dtype=np.uint8)\n            \n            # Đếm frequency của hex bằng C-backend\n            byte_counts = np.bincount(int_data, minlength=256)\n            for i, count in enumerate(byte_counts):\n                features[f'byte_{i:02X}'] = count\n                \n            height = len(int_data) // width\n            img_array = int_data[:height * width].reshape((height, width))\n            \n            # Save raw aspect ratio, DO NOT resize.\n            image = Image.fromarray(img_array)\n            image.save(os.path.join(IMAGE_DIR, f\"{fid}.png\"), format=\"PNG\")\n            \n            os.remove(byte_path) \n\n        if os.path.exists(asm_path):\n            with open(asm_path, 'r', encoding='latin1') as f:\n                asm_content = f.read().lower() # lower case for robust opcode matching\n            \n            features['asm_size'] = len(asm_content)\n            \n            for seg in SEGMENTS:\n                features[f'seg_{seg.replace(\".\", \"\")}'] = asm_content.count(seg + ':')\n                \n            # EXTRACT CORE EXECUTION OPCODES\n            for op in OPCODES:\n                # Add spaces to avoid partial matches (e.g., 'jmp' inside 'ejmp')\n                features[f'op_{op}'] = asm_content.count(f' {op} ')\n            \n            os.remove(asm_path) \n            \n    except Exception as e:\n        features['error'] = str(e)\n        \n    return features","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3. Concurrency Orchestration & Execution\nThe final block integrates the worker nodes into a robust data pipeline. \n* It uses the `7z` command-line tool to surgically extract precisely 400 files directly into `/dev/shm`.\n* A `ProcessPoolExecutor` deploys all available CPU cores to digest the batch.\n* **Checkpoint Resilience:** If the 12-hour Kaggle runtime limit kills the kernel, the engine automatically detects existing `.csv` checkpoints and resumes execution from the exact point of failure.","metadata":{}},{"cell_type":"code","source":"# --- 3. ORCHESTRATION & EXECUTION ---\ndef process_batch(id_list):\n    os.makedirs(BATCH_DIR, exist_ok=True)\n    \n    with open('batch_list.txt', 'w') as f:\n        for fid in id_list:\n            f.write(f\"{MODE}/{fid}.bytes\\n\")\n            f.write(f\"{MODE}/{fid}.asm\\n\")\n            \n    subprocess.run(['7z', 'e', ARCHIVE_PATH, '@batch_list.txt', f'-o{BATCH_DIR}', '-y'], \n                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)\n    \n    batch_data = []\n    with ProcessPoolExecutor(max_workers=MAX_WORKERS) as executor:\n        futures = {executor.submit(process_single_file, fid): fid for fid in id_list}\n        for future in as_completed(futures):\n            batch_data.append(future.result())\n            \n    if os.path.exists('batch_list.txt'): os.remove('batch_list.txt')\n    return batch_data\n\nexisting_csvs = [f for f in os.listdir(WORKING_DIR) if f.startswith(f'{MODE}_tabular_features_batch_')]\n\nif existing_csvs:\n    # --- SMART CHECKPOINT LOGIC ---\n    print(\"Scanning existing checkpoints...\")\n    processed_ids = set()\n    done_batches = []\n    \n    for f in existing_csvs:\n        done_batches.append(int(f.split('_')[-1].split('.')[0]))\n        df_tmp = pd.read_csv(os.path.join(WORKING_DIR, f))\n        processed_ids.update(df_tmp['Id'].tolist())\n        \n    target_ids = [fid for fid in target_ids if fid not in processed_ids]\n    next_batch_num = max(done_batches) + 1\n    \n    print(f\"Smart Checkpoint: Found {len(processed_ids)} files already processed safely.\")\n    print(f\"Remaining files to process: {len(target_ids)}. Resuming and saving to batch file {next_batch_num}...\")\nelse:\n    next_batch_num = 1\n\nprint(f\"INITIATING PHASE 1 PIPELINE [TARGET: {MODE.upper()}]\")\nprint(\"-\" * 50)\n\nfor i in range(0, len(target_ids), BATCH_SIZE):\n    batch_ids = target_ids[i : i + BATCH_SIZE]\n    batch_num = next_batch_num + (i // BATCH_SIZE)\n    \n    print(f\"[Batch {batch_num:03d}] Processing {len(batch_ids)} signatures...\")\n    \n    results = process_batch(batch_ids)\n    df_batch = pd.DataFrame(results)\n    \n    if 'error' in df_batch.columns:\n        error_df = df_batch[df_batch['error'].notna()]\n        if not error_df.empty:\n            print(f\"  --> [WARNING] Isolated {len(error_df)} corrupted files.\")\n        df_batch = df_batch.drop(columns=['error'])\n        \n    df_batch.to_csv(f'{MODE}_tabular_features_batch_{batch_num}.csv', index=False)\n\nprint(\"-\" * 50)\nprint(\"PHASE 1 COMPLETE.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}