import lief
import numpy as np
import json

def extract_features(pe_path):
    """Extract static features from a PE file using lief."""
    features = {}

    try:
        binary = lief.parse(pe_path)
    except Exception as e:
        return None, f"Failed to parse PE file: {e}"

    if binary is None:
        return None, "Not a valid PE file"

    # Header features
    features['has_debug'] = int(binary.has_debug)
    features['has_relocations'] = int(binary.has_relocations)
    features['has_resources'] = int(binary.has_resources)
    features['has_signature'] = int(binary.has_signatures)
    features['has_tls'] = int(binary.has_tls)
    features['has_configuration'] = int(binary.has_configuration)

    # Section features
    features['num_sections'] = len(binary.sections)
    section_sizes = [s.size for s in binary.sections]
    section_entropies = [s.entropy for s in binary.sections]
    features['section_mean_entropy'] = float(np.mean(section_entropies)) if section_entropies else 0.0
    features['section_max_entropy'] = float(np.max(section_entropies)) if section_entropies else 0.0
    features['section_mean_size'] = float(np.mean(section_sizes)) if section_sizes else 0.0

    # Import features
    if binary.has_imports:
        features['num_imports'] = sum(len(lib.entries) for lib in binary.imports)
        features['num_import_libs'] = len(binary.imports)
    else:
        features['num_imports'] = 0
        features['num_import_libs'] = 0

    # Export features
    if binary.has_exports:
        features['num_exports'] = len(binary.get_export().entries)
    else:
        features['num_exports'] = 0

    # Optional header
    oh = binary.optional_header
    features['virtual_size'] = oh.sizeof_image
    features['num_rva_and_sizes'] = oh.numberof_rva_and_size

    return features, None


if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("Usage: python extract_features.py <path_to_exe>")
        sys.exit(1)

    path = sys.argv[1]
    feats, err = extract_features(path)
    if err:
        print(f"Error: {err}")
    else:
        print(json.dumps(feats, indent=2))