""" Data Preprocessing Pipeline for Cyber-LLM Handles cleaning, tokenization, and preparation of cybersecurity training data. """ import os import json import logging import argparse from pathlib import Path from typing import Dict, List, Any, Optional, Tuple import re from dataclasses import dataclass try: from transformers import AutoTokenizer import numpy as np from sklearn.model_selection import train_test_split import dvc.api except ImportError: print("Required packages not installed. Run: pip install transformers scikit-learn dvc") # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) @dataclass class ProcessingConfig: """Configuration for data preprocessing.""" max_sequence_length: int = 2048 min_sequence_length: int = 128 overlap_ratio: float = 0.1 validation_split: float = 0.15 test_split: float = 0.1 tokenizer_name: str = "microsoft/DialoGPT-medium" special_tokens: Dict[str, str] = None def __post_init__(self): if self.special_tokens is None: self.special_tokens = { "recon_token": "<|RECON|>", "c2_token": "<|C2|>", "postexploit_token": "<|POSTEXPLOIT|>", "opsec_token": "<|OPSEC|>", "explain_token": "<|EXPLAIN|>", "safety_token": "<|SAFETY|>" } class CyberDataPreprocessor: """Advanced data preprocessor for cybersecurity domain.""" def __init__(self, config: ProcessingConfig): self.config = config self.tokenizer = None self.document_stats = {} self._initialize_tokenizer() def _initialize_tokenizer(self): """Initialize tokenizer with special tokens.""" logger.info(f"Loading tokenizer: {self.config.tokenizer_name}") self.tokenizer = AutoTokenizer.from_pretrained( self.config.tokenizer_name, trust_remote_code=True ) # Add special tokens special_tokens_dict = { 'additional_special_tokens': list(self.config.special_tokens.values()) } num_added_toks = self.tokenizer.add_special_tokens(special_tokens_dict) logger.info(f"Added {num_added_toks} special tokens") # Set pad token if not exists if self.tokenizer.pad_token is None: self.tokenizer.pad_token = self.tokenizer.eos_token def load_raw_data(self, data_dir: Path) -> List[Dict]: """Load and parse raw JSON data files.""" logger.info(f"Loading data from: {data_dir}") json_files = list(data_dir.glob('**/*.json')) documents = [] for json_file in json_files: if json_file.name in ['conversion_report.json', 'metadata.json']: continue try: with open(json_file, 'r', encoding='utf-8') as f: data = json.load(f) # Extract relevant fields if 'content' in data and 'metadata' in data: doc = { 'id': json_file.stem, 'content': data['content'], 'source_type': data['metadata'].get('source_type', 'unknown'), 'filename': data['metadata'].get('filename', ''), 'source_file': str(json_file) } documents.append(doc) except Exception as e: logger.error(f"Error loading {json_file}: {str(e)}") logger.info(f"Loaded {len(documents)} documents") return documents def clean_and_structure_text(self, documents: List[Dict]) -> List[Dict]: """Clean and structure text content.""" logger.info("Cleaning and structuring text content") cleaned_documents = [] for doc in documents: try: content = doc['content'] # Clean content cleaned_content = self._clean_text(content) # Structure based on document type structured_content = self._structure_by_type( cleaned_content, doc['source_type'] ) # Add special tokens based on content type tagged_content = self._add_domain_tags( structured_content, doc['source_type'] ) doc['cleaned_content'] = cleaned_content doc['structured_content'] = structured_content doc['tagged_content'] = tagged_content cleaned_documents.append(doc) except Exception as e: logger.error(f"Error processing document {doc['id']}: {str(e)}") continue logger.info(f"Cleaned {len(cleaned_documents)} documents") return cleaned_documents def _clean_text(self, text: str) -> str: """Clean raw text content.""" # Remove excessive whitespace text = re.sub(r'\s+', ' ', text) # Remove page numbers and headers/footers text = re.sub(r'Page \d+', '', text) text = re.sub(r'\n{3,}', '\n\n', text) # Fix common OCR errors text = re.sub(r'([a-z])([A-Z])', r'\1 \2', text) # Missing spaces text = re.sub(r'(\w)(\d)', r'\1 \2', text) # Word-number separation # Normalize quotes and dashes text = text.replace('"', '"').replace('"', '"') text = text.replace('–', '-').replace('—', '-') # Remove artifacts text = re.sub(r'^\s*[•\-\*]\s*', '', text, flags=re.MULTILINE) return text.strip() def _structure_by_type(self, content: str, doc_type: str) -> str: """Structure content based on document type.""" if doc_type == 'mitre_attack': return self._structure_mitre_content(content) elif doc_type == 'apt_report': return self._structure_apt_report(content) elif doc_type == 'opsec_guide': return self._structure_opsec_content(content) elif doc_type == 'malware_analysis': return self._structure_malware_content(content) else: return content def _structure_mitre_content(self, content: str) -> str: """Structure MITRE ATT&CK content.""" # Find and structure tactics/techniques sections = [] # Look for technique patterns technique_pattern = r'(T\d{4}(?:\.\d{3})?)\s*[-:]\s*([^\n]+)' techniques = re.findall(technique_pattern, content) if techniques: sections.append("MITRE ATT&CK Techniques:") for tech_id, tech_name in techniques: sections.append(f"- {tech_id}: {tech_name}") # Look for tactic sections tactic_pattern = r'(Initial Access|Execution|Persistence|Privilege Escalation|Defense Evasion|Credential Access|Discovery|Lateral Movement|Collection|Exfiltration|Command and Control|Impact)' current_section = None for line in content.split('\n'): tactic_match = re.search(tactic_pattern, line, re.IGNORECASE) if tactic_match: current_section = tactic_match.group(1) sections.append(f"\n{current_section}:") elif current_section and line.strip(): sections.append(f" {line.strip()}") return '\n'.join(sections) if sections else content def _structure_apt_report(self, content: str) -> str: """Structure APT report content.""" sections = [] # Extract IOCs ip_pattern = r'\b(?:\d{1,3}\.){3}\d{1,3}\b' domain_pattern = r'\b[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?(?:\.[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)+\b' hash_pattern = r'\b[a-fA-F0-9]{32,64}\b' ips = re.findall(ip_pattern, content) domains = re.findall(domain_pattern, content) hashes = re.findall(hash_pattern, content) if ips or domains or hashes: sections.append("Indicators of Compromise (IOCs):") if ips: sections.append(f"IPs: {', '.join(set(ips[:10]))}") # Limit to first 10 if domains: sections.append(f"Domains: {', '.join(set(domains[:10]))}") if hashes: sections.append(f"Hashes: {', '.join(set(hashes[:5]))}") # Extract TTPs ttp_keywords = ['technique', 'tactic', 'procedure', 'method', 'attack'] ttp_lines = [] for line in content.split('\n'): if any(keyword in line.lower() for keyword in ttp_keywords): ttp_lines.append(line.strip()) if ttp_lines: sections.append("\nTTPs (Tactics, Techniques, Procedures):") sections.extend(f"- {line}" for line in ttp_lines[:10]) sections.append(f"\nFull Report:\n{content}") return '\n'.join(sections) def _structure_opsec_content(self, content: str) -> str: """Structure OPSEC guide content.""" sections = [] # Look for OPSEC principles/rules opsec_keywords = ['stealth', 'detection', 'evasion', 'anonymity', 'operational security'] opsec_lines = [] for line in content.split('\n'): if any(keyword in line.lower() for keyword in opsec_keywords): opsec_lines.append(line.strip()) if opsec_lines: sections.append("OPSEC Guidelines:") sections.extend(f"- {line}" for line in opsec_lines[:15]) sections.append(f"\nDetailed Content:\n{content}") return '\n'.join(sections) def _structure_malware_content(self, content: str) -> str: """Structure malware analysis content.""" sections = [] # Extract analysis sections analysis_sections = ['summary', 'behavior', 'network', 'filesystem', 'registry'] for section in analysis_sections: pattern = rf'{section}[:\s]+(.*?)(?=\n[a-z]+:|$)' match = re.search(pattern, content, re.IGNORECASE | re.DOTALL) if match: sections.append(f"{section.title()}: {match.group(1).strip()}") return '\n'.join(sections) if sections else content def _add_domain_tags(self, content: str, doc_type: str) -> str: """Add domain-specific special tokens.""" tag_mapping = { 'mitre_attack': self.config.special_tokens['recon_token'], 'apt_report': self.config.special_tokens['postexploit_token'], 'opsec_guide': self.config.special_tokens['opsec_token'], 'malware_analysis': self.config.special_tokens['postexploit_token'] } tag = tag_mapping.get(doc_type, '') if tag: return f"{tag} {content}" return content def create_training_sequences(self, documents: List[Dict]) -> List[Dict]: """Create training sequences with proper tokenization.""" logger.info("Creating training sequences") sequences = [] for doc in documents: content = doc['tagged_content'] # Tokenize content tokens = self.tokenizer.encode(content, add_special_tokens=True) # Create overlapping sequences seq_length = self.config.max_sequence_length overlap = int(seq_length * self.config.overlap_ratio) for i in range(0, len(tokens), seq_length - overlap): seq_tokens = tokens[i:i + seq_length] # Skip sequences that are too short if len(seq_tokens) < self.config.min_sequence_length: continue # Pad sequence if necessary if len(seq_tokens) < seq_length: seq_tokens.extend([self.tokenizer.pad_token_id] * (seq_length - len(seq_tokens))) sequence = { 'input_ids': seq_tokens, 'attention_mask': [1 if token != self.tokenizer.pad_token_id else 0 for token in seq_tokens], 'labels': seq_tokens.copy(), # For language modeling 'source_type': doc['source_type'], 'document_id': doc['id'], 'sequence_index': i // (seq_length - overlap) } sequences.append(sequence) logger.info(f"Created {len(sequences)} training sequences") return sequences def split_data(self, sequences: List[Dict]) -> Tuple[List[Dict], List[Dict], List[Dict]]: """Split data into train/validation/test sets.""" logger.info("Splitting data into train/validation/test sets") # Group sequences by document to ensure proper splitting doc_sequences = {} for seq in sequences: doc_id = seq['document_id'] if doc_id not in doc_sequences: doc_sequences[doc_id] = [] doc_sequences[doc_id].append(seq) # Split document IDs doc_ids = list(doc_sequences.keys()) # First split: train + temp vs test train_temp_ids, test_ids = train_test_split( doc_ids, test_size=self.config.test_split, random_state=42, shuffle=True ) # Second split: train vs validation val_size = self.config.validation_split / (1 - self.config.test_split) train_ids, val_ids = train_test_split( train_temp_ids, test_size=val_size, random_state=42, shuffle=True ) # Collect sequences for each split train_sequences = [] val_sequences = [] test_sequences = [] for doc_id, doc_seqs in doc_sequences.items(): if doc_id in train_ids: train_sequences.extend(doc_seqs) elif doc_id in val_ids: val_sequences.extend(doc_seqs) else: # test_ids test_sequences.extend(doc_seqs) logger.info(f"Split: {len(train_sequences)} train, {len(val_sequences)} val, {len(test_sequences)} test") return train_sequences, val_sequences, test_sequences def save_processed_data(self, train_data: List[Dict], val_data: List[Dict], test_data: List[Dict], output_dir: Path): """Save processed data to files.""" output_dir.mkdir(parents=True, exist_ok=True) # Save datasets datasets = { 'train': train_data, 'validation': val_data, 'test': test_data } for split_name, data in datasets.items(): output_file = output_dir / f'{split_name}.json' logger.info(f"Saving {len(data)} {split_name} sequences to {output_file}") with open(output_file, 'w', encoding='utf-8') as f: json.dump(data, f, indent=2, ensure_ascii=False) # Save tokenizer tokenizer_dir = output_dir / 'tokenizer' self.tokenizer.save_pretrained(tokenizer_dir) # Save preprocessing metadata metadata = { 'config': self.config.__dict__, 'splits': { 'train_size': len(train_data), 'validation_size': len(val_data), 'test_size': len(test_data) }, 'tokenizer_info': { 'vocab_size': self.tokenizer.vocab_size, 'model_max_length': self.tokenizer.model_max_length, 'special_tokens': self.config.special_tokens } } metadata_file = output_dir / 'preprocessing_metadata.json' with open(metadata_file, 'w', encoding='utf-8') as f: json.dump(metadata, f, indent=2, ensure_ascii=False) logger.info(f"Preprocessing complete. Data saved to {output_dir}") def main(): parser = argparse.ArgumentParser(description='Preprocess cybersecurity data for Cyber-LLM') parser.add_argument('--input', required=True, help='Input directory with raw JSON files') parser.add_argument('--output', required=True, help='Output directory for processed data') parser.add_argument('--config', help='Configuration file path') parser.add_argument('--max-length', type=int, default=2048, help='Maximum sequence length') parser.add_argument('--tokenizer', default='microsoft/DialoGPT-medium', help='Tokenizer model name') args = parser.parse_args() # Create configuration config = ProcessingConfig( max_sequence_length=args.max_length, tokenizer_name=args.tokenizer ) # Initialize preprocessor preprocessor = CyberDataPreprocessor(config) # Load and process data input_dir = Path(args.input) output_dir = Path(args.output) # Load raw data documents = preprocessor.load_raw_data(input_dir) # Clean and structure cleaned_documents = preprocessor.clean_and_structure_text(documents) # Create training sequences sequences = preprocessor.create_training_sequences(cleaned_documents) # Split data train_data, val_data, test_data = preprocessor.split_data(sequences) # Save processed data preprocessor.save_processed_data(train_data, val_data, test_data, output_dir) logger.info("Data preprocessing completed successfully!") if __name__ == '__main__': main()