pyrograph 0.1.0

use crate::ir::TaintGraph;
use crate::labels::{label_node, LabelSet, SanitizerDef, SinkDef, SourceDef};

pub fn default_label_set() -> LabelSet {
    LabelSet {
        sources: vec![
            SourceDef {
                id: "go-env-getenv".into(),
                pattern: "os.Getenv".into(),
                category: "credential".into(),
            },
            SourceDef {
                id: "go-env-environ".into(),
                pattern: "os.Environ".into(),
                category: "system".into(),
            },
            SourceDef {
                id: "go-argv".into(),
                pattern: "os.Args".into(),
                category: "cli".into(),
            },
            SourceDef {
                id: "go-flag-args".into(),
                pattern: "flag.Args".into(),
                category: "cli".into(),
            },
            SourceDef {
                id: "go-readfile-ioutil".into(),
                pattern: "ioutil.ReadFile".into(),
                category: "file".into(),
            },
            SourceDef {
                id: "go-readfile-os".into(),
                pattern: "os.ReadFile".into(),
                category: "file".into(),
            },
            SourceDef {
                id: "go-http-form-value".into(),
                pattern: "http.Request.FormValue".into(),
                category: "http".into(),
            },
            SourceDef {
                id: "go-url-query".into(),
                pattern: "URL.Query".into(),
                category: "http".into(),
            },
            SourceDef {
                id: "go-net-conn".into(),
                pattern: "net.Conn".into(),
                category: "network-input".into(),
            },
            SourceDef { id: "go-ssh-key".into(), pattern: ".ssh/id_rsa".into(), category: "sensitive-file".into() },
            SourceDef { id: "go-aws-creds".into(), pattern: ".aws/credentials".into(), category: "sensitive-file".into() },
            SourceDef { id: "go-kube-config".into(), pattern: ".kube/config".into(), category: "sensitive-file".into() },
            SourceDef { id: "go-home-dir".into(), pattern: "os.UserHomeDir".into(), category: "system".into() },
            SourceDef { id: "go-hostname".into(), pattern: "os.Hostname".into(), category: "system".into() },
            // Shell/URL sources
            SourceDef { id: "go-curl".into(), pattern: "curl".into(), category: "shell".into() },
            SourceDef { id: "go-wget".into(), pattern: "wget".into(), category: "shell".into() },
            SourceDef { id: "go-bin-sh".into(), pattern: "/bin/sh".into(), category: "shell".into() },
            SourceDef { id: "go-bin-bash".into(), pattern: "/bin/bash".into(), category: "shell".into() },
            SourceDef { id: "go-http-url".into(), pattern: "http://".into(), category: "network-input".into() },
            SourceDef { id: "go-https-url".into(), pattern: "https://".into(), category: "network-input".into() },
        ],
        sinks: vec![
            SinkDef {
                id: "go-exec-command-run".into(),
                pattern: "exec.Command.Run".into(),
                category: "exec".into(),
            },
            SinkDef {
                id: "go-exec-command-output".into(),
                pattern: "exec.Command.Output".into(),
                category: "exec".into(),
            },
            SinkDef {
                id: "go-exec-command-combined".into(),
                pattern: "exec.Command.CombinedOutput".into(),
                category: "exec".into(),
            },
            SinkDef {
                id: "go-start-process".into(),
                pattern: "os.StartProcess".into(),
                category: "exec".into(),
            },
            SinkDef {
                id: "go-net-dial".into(),
                pattern: "net.Dial".into(),
                category: "network".into(),
            },
            SinkDef {
                id: "go-http-get".into(),
                pattern: "http.Get".into(),
                category: "network".into(),
            },
            SinkDef {
                id: "go-http-post".into(),
                pattern: "http.Post".into(),
                category: "network".into(),
            },
            SinkDef {
                id: "go-os-create".into(),
                pattern: "os.Create".into(),
                category: "file".into(),
            },
            SinkDef {
                id: "go-os-writefile".into(),
                pattern: "os.WriteFile".into(),
                category: "file".into(),
            },
            SinkDef {
                id: "go-reflect-call".into(),
                pattern: "reflect.ValueOf.Call".into(),
                category: "exec".into(),
            },
            SinkDef { id: "go-exec-command".into(), pattern: "exec.Command".into(), category: "exec".into() },
            SinkDef { id: "go-cgo-system".into(), pattern: "C.system".into(), category: "exec".into() },
            SinkDef { id: "go-cgo-popen".into(), pattern: "C.popen".into(), category: "exec".into() },
            SinkDef { id: "go-net-dialtcp".into(), pattern: "net.DialTCP".into(), category: "network".into() },
            SinkDef { id: "go-http-newrequest".into(), pattern: "http.NewRequest".into(), category: "network".into() },
            SinkDef { id: "go-ioutil-writefile".into(), pattern: "ioutil.WriteFile".into(), category: "file".into() },
        ],
        sanitizers: vec![
            SanitizerDef {
                id: "go-strconv-atoi".into(),
                pattern: "strconv.Atoi".into(),
            },
            SanitizerDef {
                id: "go-strconv-parseint".into(),
                pattern: "strconv.ParseInt".into(),
            },
            SanitizerDef {
                id: "go-filepath-base".into(),
                pattern: "filepath.Base".into(),
            },
            SanitizerDef {
                id: "go-html-escape".into(),
                pattern: "html.EscapeString".into(),
            },
        ],
    }
}

pub(super) fn apply_labels(graph: &mut TaintGraph, label_set: &LabelSet) {
    for id in 0..graph.node_count() as u32 {
        if let Some(node) = graph.node(id) {
            let label = label_node(label_set, &node.name)
                .or_else(|| node.alias.as_deref().and_then(|alias| label_node(label_set, alias)));

            if let Some(label) = label {
                let Some(node_mut) = graph.node_mut(id) else {
                    continue;
                };

                node_mut.label = Some(label);
            }
        }
    }
}