void vulnerable(char *input) {
    char buffer[64];
    strcpy(buffer, input);  // Dangerous!
}

CPG = AST + CFG + PDG

# Install Joern (macOS)
brew install joern

# Generate CPG (using OpenSSL as example)
joern-parse /path/to/openssl --language c --output openssl.cpg

# Export to DOT format (for visualization/debugging)
joern-export openssl.cpg --repr cpg14 --format dot

// Find all strcpy call sites
val sinks = cpg.call.name("strcpy").l

// Backward data flow slicing (depth 5 steps)
val slice = sinks.flatMap { sink =>
  sink.reachableByFlows(cpg.method.parameter, 5)
}

import torch
from torch_geometric.nn import GATConv, global_mean_pool
from torch_geometric.data import DataLoader

class VulnDetector(torch.nn.Module):
    def __init__(self, in_dim=128, hidden_dim=256, heads=8):
        super().__init__()
        self.conv1 = GATConv(in_dim, hidden_dim, heads=heads)
        self.conv2 = GATConv(hidden_dim * heads, hidden_dim, heads=1)
        self.classifier = torch.nn.Linear(hidden_dim, 2)

    def forward(self, x, edge_index, batch):
        # Two GAT layers
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index).relu()
        # Graph-level pooling
        x = global_mean_pool(x, batch)
        return self.classifier(x)

# Training loop
model = VulnDetector()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

for epoch in range(100):
    for batch in train_loader:
        optimizer.zero_grad()
        out = model(batch.x, batch.edge_index, batch.batch)
        loss = F.cross_entropy(out, batch.y)
        loss.backward()
        optimizer.step()

# Focal Loss implementation
class FocalLoss(nn.Module):
    def __init__(self, gamma=2.0, alpha=0.75):
        super().__init__()
        self.gamma = gamma
        self.alpha = alpha

    def forward(self, pred, target):
        ce_loss = F.cross_entropy(pred, target, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
        return focal_loss.mean()

// Find SQL injection
from Call call, DataFlow::PathNode source, DataFlow::PathNode sink
where
  call.getTarget().hasName("query") and
  source.isSource() and
  sink.isSink() and
  DataFlow::localFlow(source, sink)
select sink, "SQL injection from $@", source, "user input"

"""
GGNN-based Vulnerability Detector
Based on DEVIGN (NeurIPS 2019) architecture
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import MessagePassing, global_mean_pool

class GGNNLayer(MessagePassing):
    """Single GGNN layer with GRU-style gating"""

    def __init__(self, hidden_dim: int, num_edge_types: int = 4):
        super().__init__(aggr='add')  # Sum aggregation
        self.hidden_dim = hidden_dim
        self.num_edge_types = num_edge_types

        # Edge-specific transformation matrices
        self.edge_mlps = nn.ModuleList([
            nn.Linear(hidden_dim, hidden_dim, bias=False)
            for _ in range(num_edge_types)
        ])

        # GRU gate
        self.gru = nn.GRUCell(hidden_dim, hidden_dim)

    def forward(self, x, edge_index, edge_type):
        # Message passing for each edge type
        aggregated = torch.zeros_like(x)
        for etype in range(self.num_edge_types):
            mask = edge_type == etype
            if mask.sum() > 0:
                edge_idx = edge_index[:, mask]
                msg = self.edge_mlps[etype](x)
                aggregated += self.propagate(edge_idx, x=msg)

# GRU update
        return self.gru(aggregated, x)

    def message(self, x_j):
        return x_j


class DevignModel(nn.Module):
    """
    DEVIGN: Effective Vulnerability Identification
    Architecture: GGNN backbone + Conv1D readout
    """

    def __init__(
        self,
        input_dim: int = 128,
        hidden_dim: int = 256,
        num_layers: int = 6,
        num_edge_types: int = 4,  # AST, CFG, DFG, CDG
        num_classes: int = 2,
        dropout: float = 0.3
    ):
        super().__init__()

        # Input projection
        self.input_proj = nn.Linear(input_dim, hidden_dim)

        # GGNN layers
        self.ggnn_layers = nn.ModuleList([
            GGNNLayer(hidden_dim, num_edge_types)
            for _ in range(num_layers)
        ])

# Conv1D for sequence modeling
        self.conv1 = nn.Conv1d(hidden_dim, hidden_dim, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(hidden_dim, hidden_dim, kernel_size=1)

        # Classifier
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, num_classes)
        )

    def forward(self, x, edge_index, edge_type, batch):
        # Project input features
        h = self.input_proj(x)

        # GGNN message passing
        for layer in self.ggnn_layers:
            h = layer(h, edge_index, edge_type)

        # Conv1D processing (requires reshaping)
        h = h.unsqueeze(0).permute(0, 2, 1)  # [1, hidden, nodes]
        h = F.relu(self.conv1(h))
        h = self.conv2(h)
        h = h.permute(0, 2, 1).squeeze(0)  # [nodes, hidden]

        # Graph-level pooling
        graph_emb = global_mean_pool(h, batch)

        return self.classifier(graph_emb)


# Usage example
if __name__ == "__main__":
    model = DevignModel(input_dim=128, hidden_dim=256, num_layers=6)

    # Simulated input (100 nodes, batch of 4 graphs)
    x = torch.randn(100, 128)
    edge_index = torch.randint(0, 100, (2, 300))
    edge_type = torch.randint(0, 4, (300,))
    batch = torch.repeat_interleave(torch.arange(4), 25)

    output = model(x, edge_index, edge_type, batch)
    print(f"Output shape: {output.shape}")  # [4, 2]

"""
CodeBERT-based node feature initialization
Using Microsoft's pre-trained CodeBERT model
"""

import torch
from transformers import AutoTokenizer, AutoModel
from typing import List

class CodeBERTEmbedder:
    """Extract code embeddings using CodeBERT"""

    def __init__(
        self,
        model_name: str = "microsoft/codebert-base",
        device: str = "cuda" if torch.cuda.is_available() else "cpu",
        max_length: int = 512
    ):
        self.device = device
        self.max_length = max_length

        # Load pre-trained model
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name).to(device)
        self.model.eval()

    @torch.no_grad()
    def encode(self, code_snippets: List[str]) -> torch.Tensor:
        """
        Encode code snippets to embedding vectors

        Args:
            code_snippets: List of code strings (one per node)

        Returns:
            Tensor of shape [num_nodes, 768]
        """
        embeddings = []

        for code in code_snippets:
            # Tokenize
            inputs = self.tokenizer(
                code,
                return_tensors="pt",
                max_length=self.max_length,
                truncation=True,
                padding="max_length"
            ).to(self.device)

            # Get [CLS] token embedding
            outputs = self.model(**inputs)
            cls_emb = outputs.last_hidden_state[:, 0, :]  # [1, 768]
            embeddings.append(cls_emb)

        return torch.cat(embeddings, dim=0)

    def encode_batch(
        self,
        code_snippets: List[str],
        batch_size: int = 32
    ) -> torch.Tensor:
        """Batch encoding for large datasets"""
        all_embeddings = []

        for i in range(0, len(code_snippets), batch_size):
            batch = code_snippets[i:i + batch_size]

            inputs = self.tokenizer(
                batch,
                return_tensors="pt",
                max_length=self.max_length,
                truncation=True,
                padding=True
            ).to(self.device)

            with torch.no_grad():
                outputs = self.model(**inputs)
                cls_emb = outputs.last_hidden_state[:, 0, :]
                all_embeddings.append(cls_emb.cpu())

        return torch.cat(all_embeddings, dim=0)


# Usage example
if __name__ == "__main__":
    embedder = CodeBERTEmbedder()

    code_samples = [
        "void foo() { int x = 0; }",
        "char* strcpy(char* dest, const char* src);",
        "if (ptr != NULL) { free(ptr); }"
    ]

    embeddings = embedder.encode(code_samples)
    print(f"Embeddings shape: {embeddings.shape}")  # [3, 768]

"""
CPG Graph Slicing Implementation
Based on DeepWukong (TOSEM 2021) methodology
"""

from typing import List, Dict, Set, Tuple
from dataclasses import dataclass
import networkx as nx

@dataclass
class CPGNode:
    """CPG node representation"""
    id: int
    type: str  # 'call', 'identifier', 'literal', etc.
    code: str
    line_number: int

@dataclass
class CPGEdge:
    """CPG edge representation"""
    src: int
    dst: int
    type: str  # 'AST', 'CFG', 'DFG', 'CDG'

class CPGSlicer:
    """Backward slicing from sensitive sinks"""

# Sensitive sink functions by vulnerability type
    SINKS = {
        'buffer_overflow': ['strcpy', 'strcat', 'sprintf', 'gets', 'memcpy'],
        'command_injection': ['system', 'popen', 'exec', 'execve', 'execl'],
        'format_string': ['printf', 'fprintf', 'sprintf', 'snprintf'],
        'memory_leak': ['malloc', 'calloc', 'realloc', 'strdup'],
        'use_after_free': ['free'],  # Track pointers after free
        'sql_injection': ['mysql_query', 'sqlite3_exec', 'PQexec'],
    }

    def __init__(
        self,
        nodes: List[CPGNode],
        edges: List[CPGEdge],
        max_depth: int = 5
    ):
        self.nodes = {n.id: n for n in nodes}
        self.max_depth = max_depth

        # Build graph
        self.graph = nx.DiGraph()
        for node in nodes:
            self.graph.add_node(node.id, data=node)

        for edge in edges:
            self.graph.add_edge(edge.src, edge.dst, type=edge.type)

        # Build reverse graph for backward slicing
        self.reverse_graph = self.graph.reverse()

    def find_sinks(self, vuln_type: str = 'all') -> List[int]:
        """Find all sink nodes in the graph"""
        sinks = []

        if vuln_type == 'all':
            sink_names = set()
            for names in self.SINKS.values():
                sink_names.update(names)
        else:
            sink_names = set(self.SINKS.get(vuln_type, []))

        for node_id, node in self.nodes.items():
            if node.type == 'call':
                # Check if function name matches a sink
                for sink_name in sink_names:
                    if sink_name in node.code:
                        sinks.append(node_id)
                        break

        return sinks

    def backward_slice(
        self,
        sink_id: int,
        edge_types: Set[str] = {'DFG', 'CDG'}
    ) -> Set[int]:
        """
        Perform backward slicing from a sink node

        Args:
            sink_id: ID of the sink node
            edge_types: Edge types to follow (DFG for data flow, CDG for control)

        Returns:
            Set of node IDs in the slice
        """
        visited = set()
        queue = [(sink_id, 0)]  # (node_id, depth)

        while queue:
            node_id, depth = queue.pop(0)

            if node_id in visited or depth > self.max_depth:
                continue

            visited.add(node_id)

            # Traverse backward edges
            for pred in self.reverse_graph.predecessors(node_id):
                edge_data = self.graph[pred][node_id]
                if edge_data['type'] in edge_types:
                    queue.append((pred, depth + 1))

        return visited

    def slice_for_vulnerability(
        self,
        vuln_type: str = 'all'
    ) -> List[Tuple[int, Set[int]]]:
        """
        Generate slices for all sinks of a given vulnerability type

        Returns:
            List of (sink_id, slice_node_set) tuples
        """
        sinks = self.find_sinks(vuln_type)
        slices = []

        for sink_id in sinks:
            slice_nodes = self.backward_slice(sink_id)
            slices.append((sink_id, slice_nodes))

        return slices

    def extract_subgraph(
        self,
        node_ids: Set[int]
    ) -> Tuple[List[CPGNode], List[CPGEdge]]:
        """Extract subgraph containing only specified nodes"""
        nodes = [self.nodes[nid] for nid in node_ids if nid in self.nodes]
        edges = []

        for src, dst, data in self.graph.edges(data=True):
            if src in node_ids and dst in node_ids:
                edges.append(CPGEdge(src, dst, data['type']))

        return nodes, edges


# Joern integration helper
def joern_query_to_slice(joern_result: str) -> List[int]:
    """
    Parse Joern query results to get slice nodes

    Joern query example:
    cpg.call.name("strcpy").reachableByFlows(cpg.method.parameter, 5).id.l
    """
    # Parse Joern output format
    import re
    ids = re.findall(r'\d+', joern_result)
    return [int(i) for i in ids]


# Usage example
if __name__ == "__main__":
    # Create sample CPG
    nodes = [
        CPGNode(0, 'param', 'char* input', 1),
        CPGNode(1, 'identifier', 'input', 2),
        CPGNode(2, 'call', 'strlen(input)', 2),
        CPGNode(3, 'identifier', 'buffer', 3),
        CPGNode(4, 'call', 'strcpy(buffer, input)', 4),  # SINK
        CPGNode(5, 'return', 'return 0', 5),
    ]

    edges = [
        CPGEdge(0, 1, 'DFG'),  # input def -> use
        CPGEdge(1, 2, 'DFG'),  # input -> strlen
        CPGEdge(1, 4, 'DFG'),  # input -> strcpy (TAINT)
        CPGEdge(3, 4, 'DFG'),  # buffer -> strcpy
        CPGEdge(2, 4, 'CDG'),  # strlen -> strcpy (control)
        CPGEdge(4, 5, 'CFG'),  # strcpy -> return
    ]

    slicer = CPGSlicer(nodes, edges, max_depth=5)

    # Find buffer overflow sinks
    sinks = slicer.find_sinks('buffer_overflow')
    print(f"Found sinks: {sinks}")  # [4]

    # Backward slice from sink
    slices = slicer.slice_for_vulnerability('buffer_overflow')
    for sink_id, slice_nodes in slices:
        print(f"Sink {sink_id}: slice contains {len(slice_nodes)} nodes")
        print(f"  Nodes: {slice_nodes}")

"""
Complete GNN Vulnerability Detection Training Pipeline
Includes: data loading, training loop, evaluation, checkpointing
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch_geometric.data import DataLoader, Data
from torch_geometric.nn import GATConv, global_mean_pool
from sklearn.metrics import f1_score, precision_score, recall_score, roc_auc_score
from typing import Dict, List, Tuple
import numpy as np
from tqdm import tqdm

class FocalLoss(nn.Module):
    """Focal Loss for class imbalance"""

    def __init__(self, gamma: float = 2.0, alpha: float = 0.75):
        super().__init__()
        self.gamma = gamma
        self.alpha = alpha

    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        ce_loss = F.cross_entropy(pred, target, reduction='none')
        pt = torch.exp(-ce_loss)
        focal_loss = self.alpha * (1 - pt) ** self.gamma * ce_loss
        return focal_loss.mean()


class VulnerabilityDetector(nn.Module):
    """GAT-based Vulnerability Detector with residual connections"""

    def __init__(
        self,
        input_dim: int = 768,  # CodeBERT dimension
        hidden_dim: int = 256,
        num_heads: int = 8,
        num_layers: int = 4,
        dropout: float = 0.3,
        num_classes: int = 2
    ):
        super().__init__()

        # Input projection
        self.input_proj = nn.Linear(input_dim, hidden_dim)

        # GAT layers with residual connections
        self.gat_layers = nn.ModuleList()
        self.layer_norms = nn.ModuleList()

        for i in range(num_layers):
            in_channels = hidden_dim if i == 0 else hidden_dim * num_heads
            self.gat_layers.append(
                GATConv(in_channels, hidden_dim, heads=num_heads, dropout=dropout)
            )
            self.layer_norms.append(nn.LayerNorm(hidden_dim * num_heads))

        # Final projection
        self.final_gat = GATConv(hidden_dim * num_heads, hidden_dim, heads=1, dropout=dropout)

        # Classifier
        self.classifier = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, num_classes)
        )

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, edge_index, batch):
        # Input projection
        h = self.input_proj(x)

        # First GAT layer (no residual - dimension change)
        h = self.gat_layers[0](h, edge_index)
        h = F.elu(h)
        h = self.layer_norms[0](h)
        h = self.dropout(h)

        # Remaining GAT layers with residual connections
        for i in range(1, len(self.gat_layers)):
            h_res = h
            h = self.gat_layers[i](h, edge_index)
            h = F.elu(h)
            h = self.layer_norms[i](h)
            h = self.dropout(h)
            h = h + h_res  # Residual connection

        # Final GAT layer
        h = self.final_gat(h, edge_index)
        h = F.elu(h)

        # Graph-level pooling
        graph_emb = global_mean_pool(h, batch)

        return self.classifier(graph_emb)


class VulnerabilityTrainer:
    """Training pipeline with evaluation and checkpointing"""

    def __init__(
        self,
        model: nn.Module,
        device: str = "cuda" if torch.cuda.is_available() else "cpu",
        learning_rate: float = 1e-4,
        weight_decay: float = 1e-5
    ):
        self.model = model.to(device)
        self.device = device

        # Loss and optimizer
        self.criterion = FocalLoss(gamma=2.0, alpha=0.75)
        self.optimizer = AdamW(
            model.parameters(),
            lr=learning_rate,
            weight_decay=weight_decay
        )

        # Metrics tracking
        self.best_f1 = 0.0
        self.history = {'train_loss': [], 'val_f1': []}

    def train_epoch(self, train_loader: DataLoader) -> float:
        """Train for one epoch"""
        self.model.train()
        total_loss = 0.0

        for batch in tqdm(train_loader, desc="Training"):
            batch = batch.to(self.device)

            self.optimizer.zero_grad()
            output = self.model(batch.x, batch.edge_index, batch.batch)
            loss = self.criterion(output, batch.y)

            loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
            self.optimizer.step()

            total_loss += loss.item()

        return total_loss / len(train_loader)

    @torch.no_grad()
    def evaluate(self, val_loader: DataLoader) -> Dict[str, float]:
        """Evaluate model on validation set"""
        self.model.eval()

        all_preds = []
        all_labels = []
        all_probs = []

        for batch in val_loader:
            batch = batch.to(self.device)
            output = self.model(batch.x, batch.edge_index, batch.batch)
            probs = F.softmax(output, dim=1)[:, 1]
            preds = output.argmax(dim=1)

            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(batch.y.cpu().numpy())
            all_probs.extend(probs.cpu().numpy())

        # Calculate metrics
        metrics = {
            'f1': f1_score(all_labels, all_preds),
            'precision': precision_score(all_labels, all_preds),
            'recall': recall_score(all_labels, all_preds),
            'auc': roc_auc_score(all_labels, all_probs) if len(set(all_labels)) > 1 else 0.0
        }

        return metrics

    def train(
        self,
        train_loader: DataLoader,
        val_loader: DataLoader,
        num_epochs: int = 100,
        patience: int = 10,
        checkpoint_path: str = "best_model.pt"
    ) -> Dict:
        """Full training loop with early stopping"""

        scheduler = CosineAnnealingLR(self.optimizer, T_max=num_epochs)
        no_improve = 0

        for epoch in range(num_epochs):
            # Training
            train_loss = self.train_epoch(train_loader)
            self.history['train_loss'].append(train_loss)

            # Evaluation
            metrics = self.evaluate(val_loader)
            self.history['val_f1'].append(metrics['f1'])

            # Learning rate scheduling
            scheduler.step()

            # Logging
            print(f"Epoch {epoch+1}/{num_epochs}")
            print(f"  Train Loss: {train_loss:.4f}")
            print(f"  Val F1: {metrics['f1']:.4f}, Precision: {metrics['precision']:.4f}, "
                  f"Recall: {metrics['recall']:.4f}, AUC: {metrics['auc']:.4f}")

            # Checkpointing
            if metrics['f1'] > self.best_f1:
                self.best_f1 = metrics['f1']
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': self.model.state_dict(),
                    'optimizer_state_dict': self.optimizer.state_dict(),
                    'best_f1': self.best_f1
                }, checkpoint_path)
                print(f"  Saved best model (F1: {self.best_f1:.4f})")
                no_improve = 0
            else:
                no_improve += 1

            # Early stopping
            if no_improve >= patience:
                print(f"Early stopping at epoch {epoch+1}")
                break

        return self.history


def create_synthetic_dataset(num_graphs: int = 1000) -> List[Data]:
    """Create synthetic dataset for testing"""
    dataset = []

    for _ in range(num_graphs):
        num_nodes = np.random.randint(50, 200)
        num_edges = np.random.randint(num_nodes, num_nodes * 3)

        # Random node features (simulating CodeBERT embeddings)
        x = torch.randn(num_nodes, 768)

        # Random edges
        edge_index = torch.randint(0, num_nodes, (2, num_edges))

        # Random label (0: safe, 1: vulnerable)
        y = torch.tensor([np.random.randint(0, 2)])

        dataset.append(Data(x=x, edge_index=edge_index, y=y))

    return dataset


# Main execution
if __name__ == "__main__":
    # Create synthetic dataset
    print("Creating synthetic dataset...")
    train_data = create_synthetic_dataset(800)
    val_data = create_synthetic_dataset(200)

    # Create data loaders
    train_loader = DataLoader(train_data, batch_size=32, shuffle=True)
    val_loader = DataLoader(val_data, batch_size=32)

    # Initialize model and trainer
    model = VulnerabilityDetector(
        input_dim=768,
        hidden_dim=256,
        num_heads=8,
        num_layers=4
    )

    print(f"Model parameters: {sum(p.numel() for p in model.parameters()):,}")

    trainer = VulnerabilityTrainer(model)

    # Train
    print("\nStarting training...")
    history = trainer.train(
        train_loader,
        val_loader,
        num_epochs=50,
        patience=10,
        checkpoint_path="vuln_detector_best.pt"
    )

    print(f"\nTraining complete. Best F1: {trainer.best_f1:.4f}")