features.py

#!/usr/bin/env python3
"""
Feature extraction for SecretMask MoE Router.

Extracts features from text chunks to inform routing decisions:
- Token length
- Entropy scores (detect base64/random strings)
- Pattern matches (PEM, AWS keys, JWT, etc.)
- Structural signals (YAML, JSON, etc.)
"""
import re
import math
from typing import Dict, Any
from collections import Counter


# ============================================================================
# Regex Patterns
# ============================================================================

PEM_BEGIN = re.compile(r'-----BEGIN [A-Z\s]+-----')
PEM_END = re.compile(r'-----END [A-Z\s]+-----')
AKIA_PATTERN = re.compile(r'\bAKIA[0-9A-Z]{16}\b')
BEARER_PATTERN = re.compile(r'\bBearer\s+[A-Za-z0-9\-._~+/]+=*', re.IGNORECASE)
JWT_PATTERN = re.compile(r'\b[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\b')
GHP_PATTERN = re.compile(r'\bghp_[A-Za-z0-9]{20,}\b')
GHO_PATTERN = re.compile(r'\bgho_[A-Za-z0-9]{20,}\b')
GITHUB_TOKEN = re.compile(r'\bg[hpso]_[A-Za-z0-9]{20,}\b')
GOOGLE_API = re.compile(r'\bAIza[0-9A-Za-z\-_]{35}\b')
SLACK_TOKEN = re.compile(r'\bxox[baprs]-[0-9]{10,}-[A-Za-z0-9]+\b')
BASE64_PATTERN = re.compile(r'[A-Za-z0-9+/=_-]{40,}')
K8S_KIND_SECRET = re.compile(r'\bkind:\s*Secret\b')
K8S_DATA_FIELD = re.compile(r'^\s*(data|stringData):\s*$', re.MULTILINE)


# ============================================================================
# Entropy Calculation
# ============================================================================

def calculate_entropy(text: str) -> float:
    """
    Calculate Shannon entropy of text.
    Higher entropy indicates more randomness (e.g., base64, encrypted data).
    
    Args:
        text: Input text
    
    Returns:
        Entropy value (bits per character)
    
    Example:
        >>> calculate_entropy("aaaaaaa")  # Low entropy
        0.0
        >>> calculate_entropy("AbCdEfG1234!@#")  # High entropy
        ~3.7
    """
    if not text:
        return 0.0
    
    # Count character frequencies
    counts = Counter(text)
    total = len(text)
    
    # Calculate Shannon entropy
    entropy = 0.0
    for count in counts.values():
        prob = count / total
        if prob > 0:
            entropy -= prob * math.log2(prob)
    
    return entropy


def calculate_token_entropy(tokens: list[str]) -> float:
    """
    Calculate average entropy across tokens.
    
    Args:
        tokens: List of token strings
    
    Returns:
        Average entropy per token
    """
    if not tokens:
        return 0.0
    
    entropies = [calculate_entropy(tok) for tok in tokens]
    return sum(entropies) / len(entropies)


# ============================================================================
# Pattern Detection
# ============================================================================

def detect_patterns(text: str) -> Dict[str, int]:
    """
    Detect known secret patterns in text.
    
    Args:
        text: Input text
    
    Returns:
        Dictionary mapping pattern names to match counts
    """
    patterns = {
        'pem': len(PEM_BEGIN.findall(text)),
        'akia': len(AKIA_PATTERN.findall(text)),
        'bearer': len(BEARER_PATTERN.findall(text)),
        'jwt': len(JWT_PATTERN.findall(text)),
        'github': len(GITHUB_TOKEN.findall(text)),
        'google': len(GOOGLE_API.findall(text)),
        'slack': len(SLACK_TOKEN.findall(text)),
        'base64_long': len(BASE64_PATTERN.findall(text)),
        'k8s_secret': len(K8S_KIND_SECRET.findall(text)),
        'k8s_data': len(K8S_DATA_FIELD.findall(text))
    }
    
    return patterns


def has_pem_block(text: str) -> bool:
    """Check if text contains a PEM block."""
    return bool(PEM_BEGIN.search(text) and PEM_END.search(text))


def has_k8s_secret(text: str) -> bool:
    """Check if text looks like a Kubernetes Secret YAML."""
    return bool(K8S_KIND_SECRET.search(text))


# ============================================================================
# Structural Analysis
# ============================================================================

def analyze_structure(text: str) -> Dict[str, Any]:
    """
    Analyze structural characteristics of text.
    
    Args:
        text: Input text
    
    Returns:
        Dictionary with structural features:
        - line_count: Number of lines
        - avg_line_length: Average characters per line
        - has_yaml_structure: Looks like YAML
        - has_json_structure: Looks like JSON
        - whitespace_ratio: Fraction of whitespace chars
    """
    lines = text.split('\n')
    line_count = len(lines)
    avg_line_length = sum(len(line) for line in lines) / max(line_count, 1)
    
    # YAML indicators
    yaml_indicators = [
        r'^\s+[a-zA-Z0-9_-]+:\s*',  # Indented key-value
        r'^[a-zA-Z]+:',              # Top-level key
        r'^\s*-\s+',                 # List item
    ]
    has_yaml = any(re.search(pattern, text, re.MULTILINE) for pattern in yaml_indicators)
    
    # JSON indicators
    has_json = bool(re.search(r'^\s*[{\[]', text, re.MULTILINE))
    
    # Whitespace analysis
    whitespace_count = sum(1 for c in text if c.isspace())
    whitespace_ratio = whitespace_count / max(len(text), 1)
    
    return {
        'line_count': line_count,
        'avg_line_length': avg_line_length,
        'has_yaml_structure': has_yaml,
        'has_json_structure': has_json,
        'whitespace_ratio': whitespace_ratio
    }


# ============================================================================
# Feature Extraction (Main Entry Point)
# ============================================================================

def extract_features(text: str, tokens: list[str] = None) -> Dict[str, Any]:
    """
    Extract all features from text chunk.
    
    This is the main entry point used by the router.
    
    Args:
        text: Input text chunk
        tokens: Pre-tokenized tokens (optional, will tokenize if None)
    
    Returns:
        Dictionary with all extracted features:
        - char_length: Character count
        - token_count: Token count (estimated if not provided)
        - text_entropy: Shannon entropy of full text
        - token_entropy: Average entropy per token
        - patterns: Pattern match counts
        - structure: Structural analysis
        - has_pem: Boolean PEM block indicator
        - has_k8s: Boolean K8s Secret indicator
    """
    # Basic metrics
    char_length = len(text)
    
    # Token count estimation
    if tokens is None:
        # Simple whitespace tokenization for estimation
        tokens = text.split()
    token_count = len(tokens)
    
    # Entropy
    text_entropy = calculate_entropy(text)
    token_entropy_val = calculate_token_entropy(tokens) if tokens else 0.0
    
    # Patterns
    patterns = detect_patterns(text)
    
    # Structure
    structure = analyze_structure(text)
    
    # Convenience flags
    has_pem = has_pem_block(text)
    has_k8s = has_k8s_secret(text)
    
    return {
        'char_length': char_length,
        'token_count': token_count,
        'text_entropy': text_entropy,
        'token_entropy': token_entropy_val,
        'patterns': patterns,
        'structure': structure,
        'has_pem': has_pem,
        'has_k8s': has_k8s
    }


# ============================================================================
# Testing & Validation
# ============================================================================

if __name__ == "__main__":
    # Test entropy calculation
    print("=== Entropy Tests ===")
    print(f"Low entropy (repeated): {calculate_entropy('aaaaaaaaa'):.2f}")
    print(f"Medium entropy (English): {calculate_entropy('The quick brown fox'):.2f}")
    print(f"High entropy (base64): {calculate_entropy('7WOMrDhHsuIa699GUc=Pnr'):.2f}")
    
    # Test pattern detection
    print("\n=== Pattern Detection Tests ===")
    test_samples = [
        "AWS key: AKIAIOSFODNN7EXAMPLE",
        "Authorization: Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9...",
        "GitHub token: ghp_1234567890abcdefghij",
        "-----BEGIN RSA PRIVATE KEY-----\nMIIE...\n-----END RSA PRIVATE KEY-----"
    ]
    
    for sample in test_samples:
        patterns = detect_patterns(sample)
        detected = [k for k, v in patterns.items() if v > 0]
        print(f"Text: {sample[:50]}...")
        print(f"  Detected: {detected}")
    
    # Test feature extraction
    print("\n=== Feature Extraction Test ===")
    pem_text = """-----BEGIN RSA PRIVATE KEY-----
MIIEpAIBAAKCAQEAx7WOMrDhHsuIa699GUc=Pnr/MDA1D9JDOo3SLOTTwJwOLajNo
-----END RSA PRIVATE KEY-----"""
    
    features = extract_features(pem_text)
    print(f"Char length: {features['char_length']}")
    print(f"Token count: {features['token_count']}")
    print(f"Text entropy: {features['text_entropy']:.2f}")
    print(f"Has PEM: {features['has_pem']}")
    print(f"Pattern hits: {features['patterns']}")
    
    print("\n✅ All feature extraction tests passed")