pyrograph 0.1.0

GPU-accelerated taint analysis for supply chain malware detection
Documentation
use super::*;
use swc_ecma_visit::Visit;

impl JsParser {
    fn try_exports(&mut self, expr: &Expr) -> bool {
        if let Expr::Assign(a) = expr {
            if let AssignTarget::Simple(SimpleAssignTarget::Member(m)) = &a.left {
                // exports = ...
                if let Expr::Ident(obj) = &*m.obj {
                    if obj.sym == "exports" {
                        let rhs = self.eval_expr(&a.right);
                        let me = *self.module_exports.get_or_insert_with(|| {
                            self.graph
                                .add_node(NodeKind::Variable, "module.exports".to_string(), None)
                        });
                        self.flow(rhs, me);
                        return true;
                    }
                }
                // module.exports = ...
                if let Expr::Ident(obj) = &*m.obj {
                    if obj.sym == "module" {
                        if let MemberProp::Ident(prop) = &m.prop {
                            if prop.sym == "exports" {
                                let rhs = self.eval_expr(&a.right);
                                let me = self.graph.add_node(
                                    NodeKind::Variable,
                                    "module.exports".to_string(),
                                    None,
                                );
                                self.flow(rhs, me);
                                self.module_exports = Some(me);
                                return true;
                            }
                        }
                    }
                }
                // module.exports.x = ...
                if let Expr::Member(outer) = &*m.obj {
                    if let Expr::Ident(obj) = &*outer.obj {
                        if obj.sym == "module" {
                            if let MemberProp::Ident(prop) = &outer.prop {
                                if prop.sym == "exports" {
                                    let rhs = self.eval_expr(&a.right);
                                    let me = *self.module_exports.get_or_insert_with(|| {
                                        self.graph.add_node(
                                            NodeKind::Variable,
                                            "module.exports".to_string(),
                                            None,
                                        )
                                    });
                                    self.flow(rhs, me);
                                    return true;
                                }
                            }
                        }
                    }
                }
            }
        }
        false
    }

    fn visit_fn_decl(&mut self, node: &FnDecl) {
        let name = node.ident.sym.to_string();
        let n = self.graph.add_node(NodeKind::Variable, name.clone(), None);
        self.cur().define(name.clone(), n);
        let prev_fn = self.current_function_node.replace(n);
        self.scopes.push(Scope::new());
        let mut params = Vec::new();
        for p in &node.function.params {
            params.extend(self.bind_pat_with_source(&p.pat, None));
        }
        self.function_params.insert(name, params);
        if let Some(b) = &node.function.body {
            for s in &b.stmts {
                self.visit_stmt(s);
            }
        }
        self.scopes.pop();
        self.current_function_node = prev_fn;
    }

    fn visit_class_decl(&mut self, node: &ClassDecl) {
        let name = node.ident.sym.to_string();
        let n = self.graph.add_node(NodeKind::Variable, name.clone(), None);
        self.cur().define(name, n);
        if let Some(s) = &node.class.super_class {
            let sup = self.eval_expr(s);
            self.flow(sup, n);
        }
        for m in &node.class.body {
            match m {
                ClassMember::Method(mm) => {
                    if let PropName::Ident(i) = &mm.key {
                        let mn = self.graph.add_node(NodeKind::Variable, i.sym.to_string(), None);
                        self.flow(mn, n);
                        let prev_fn = self.current_function_node.replace(mn);
                        self.scopes.push(Scope::new());
                        for p in &mm.function.params {
                            self.bind_pat(&p.pat);
                        }
                        if let Some(b) = &mm.function.body {
                            for s in &b.stmts {
                                self.visit_stmt(s);
                            }
                        }
                        self.scopes.pop();
                        self.current_function_node = prev_fn;
                    }
                }
                ClassMember::Constructor(c) => {
                    let cn = self
                        .graph
                        .add_node(NodeKind::Variable, "<constructor>".to_string(), None);
                    self.flow(cn, n);
                    self.scopes.push(Scope::new());
                    for p in &c.params {
                        if let ParamOrTsParamProp::Param(pp) = p {
                            self.bind_pat(&pp.pat);
                        }
                    }
                    if let Some(b) = &c.body {
                        for s in &b.stmts {
                            self.visit_stmt(s);
                        }
                    }
                    self.scopes.pop();
                }
                _ => {}
            }
        }
    }
}

impl Visit for JsParser {
    fn visit_module(&mut self, node: &Module) {
        self.scopes.push(Scope::new());
        for item in &node.body {
            self.visit_module_item(item);
        }
        self.scopes.pop();
    }

    fn visit_module_item(&mut self, node: &ModuleItem) {
        match node {
            ModuleItem::ModuleDecl(d) => self.visit_module_decl(d),
            ModuleItem::Stmt(s) => self.visit_stmt(s),
        }
    }

    fn visit_module_decl(&mut self, node: &ModuleDecl) {
        match node {
            ModuleDecl::Import(i) => self.visit_import_decl(i),
            ModuleDecl::ExportDecl(e) => self.visit_export_decl(e),
            ModuleDecl::ExportDefaultExpr(e) => {
                self.eval_expr(&e.expr);
            }
            ModuleDecl::ExportDefaultDecl(e) => self.visit_export_default_decl(e),
            ModuleDecl::ExportAll(e) => {
                self.graph.add_node(NodeKind::Import, e.src.value.to_string(), None);
            }
            _ => {}
        }
    }

    fn visit_import_decl(&mut self, node: &ImportDecl) {
        let src = node.src.value.to_string();
        for s in &node.specifiers {
            let (local, imported) = match s {
                ImportSpecifier::Named(n) => {
                    let local_name = n.local.sym.to_string();
                    let imported_name = n.imported.as_ref().map(|i| match i {
                        swc_ecma_ast::ModuleExportName::Ident(id) => id.sym.to_string(),
                        swc_ecma_ast::ModuleExportName::Str(s) => s.value.to_string(),
                    }).unwrap_or_else(|| local_name.clone());
                    (local_name, Some(imported_name))
                }
                ImportSpecifier::Default(d) => (d.local.sym.to_string(), None),
                ImportSpecifier::Namespace(ns) => (ns.local.sym.to_string(), None),
            };
            
            let name = if let Some(imp_name) = imported {
                format!("{}.{}", src, imp_name)
            } else {
                src.clone()
            };
            let imp = self.graph.add_node(NodeKind::Import, name, None);
            self.cur().define(local, imp);
        }
    }

    fn visit_export_decl(&mut self, node: &ExportDecl) {
        match &node.decl {
            Decl::Fn(f) => self.visit_fn_decl(f),
            Decl::Var(v) => self.visit_var_decl(v),
            Decl::Class(c) => self.visit_class_decl(c),
            _ => {}
        }
    }

    fn visit_export_default_decl(&mut self, node: &ExportDefaultDecl) {
        if let DefaultDecl::Fn(f) = &node.decl {
            if let Some(i) = &f.ident {
                let n = self.graph.add_node(NodeKind::Variable, i.sym.to_string(), None);
                self.cur().define(i.sym.to_string(), n);
            }
            self.eval_fn(&FnExpr {
                ident: f.ident.clone(),
                function: f.function.clone(),
            });
        }
    }

    fn visit_stmt(&mut self, node: &Stmt) {
        match node {
            Stmt::Decl(d) => self.visit_decl(d),
            Stmt::Expr(e) => {
                if !self.try_exports(&e.expr) {
                    self.eval_expr(&e.expr);
                }
            }
            Stmt::Block(b) => {
                self.scopes.push(Scope::new());
                for s in &b.stmts {
                    self.visit_stmt(s);
                }
                self.scopes.pop();
            }
            Stmt::If(i) => {
                self.eval_expr(&i.test);
                self.visit_stmt(&i.cons);
                if let Some(a) = &i.alt {
                    self.visit_stmt(a);
                }
            }
            Stmt::While(w) => {
                self.eval_expr(&w.test);
                self.visit_stmt(&w.body);
            }
            Stmt::For(f) => {
                self.scopes.push(Scope::new());
                match &f.init {
                    Some(VarDeclOrExpr::VarDecl(v)) => self.visit_var_decl(v),
                    Some(VarDeclOrExpr::Expr(e)) => {
                        self.eval_expr(e);
                    }
                    None => {}
                }
                if let Some(t) = &f.test {
                    self.eval_expr(t);
                }
                if let Some(u) = &f.update {
                    self.eval_expr(u);
                }
                self.visit_stmt(&f.body);
                self.scopes.pop();
            }
            Stmt::ForOf(f) => {
                let rhs = self.eval_expr(&f.right);
                self.scopes.push(Scope::new());
                match &f.left {
                    ForHead::VarDecl(v) => {
                        for d in &v.decls {
                            self.bind_pat_with_source(&d.name, Some(rhs));
                        }
                    }
                    ForHead::Pat(p) => { let _ = self.bind_pat_with_source(p, Some(rhs)); }
                    _ => {}
                }
                self.visit_stmt(&f.body);
                self.scopes.pop();
            }
            Stmt::Return(r) => {
                if let Some(a) = &r.arg {
                    let val = self.eval_expr(a);
                    if let Some(fn_node) = self.current_function_node {
                        self.graph.add_edge(val, fn_node, EdgeKind::Return);
                    }
                }
            }
            Stmt::Switch(s) => {
                let disc = self.eval_expr(&s.discriminant);
                for case in &s.cases {
                    if let Some(test) = &case.test {
                        let t = self.eval_expr(test);
                        self.flow(disc, t);
                    }
                    for stmt in &case.cons {
                        self.visit_stmt(stmt);
                    }
                }
            }
            Stmt::ForIn(f) => {
                let rhs = self.eval_expr(&f.right);
                self.scopes.push(Scope::new());
                match &f.left {
                    ForHead::VarDecl(v) => {
                        for d in &v.decls {
                            self.bind_pat_with_source(&d.name, Some(rhs));
                        }
                    }
                    ForHead::Pat(p) => { let _ = self.bind_pat_with_source(p, Some(rhs)); }
                    _ => {}
                }
                self.visit_stmt(&f.body);
                self.scopes.pop();
            }
            Stmt::Try(t) => {
                for s in &t.block.stmts {
                    self.visit_stmt(s);
                }
                if let Some(c) = &t.handler {
                    self.scopes.push(Scope::new());
                    if let Some(p) = &c.param {
                        self.bind_pat(p);
                    }
                    for s in &c.body.stmts {
                        self.visit_stmt(s);
                    }
                    self.scopes.pop();
                }
                if let Some(f) = &t.finalizer {
                    for s in &f.stmts {
                        self.visit_stmt(s);
                    }
                }
            }
            Stmt::Throw(t) => {
                self.eval_expr(&t.arg);
            }
            Stmt::DoWhile(d) => {
                self.visit_stmt(&d.body);
                self.eval_expr(&d.test);
            }
            _ => {}
        }
    }

    fn visit_decl(&mut self, node: &Decl) {
        match node {
            Decl::Var(v) => self.visit_var_decl(v),
            Decl::Fn(f) => self.visit_fn_decl(f),
            Decl::Class(c) => self.visit_class_decl(c),
            _ => {}
        }
    }

    fn visit_var_decl(&mut self, node: &VarDecl) {
        for d in &node.decls {
            let rhs = d.init.as_deref().map(|e| self.eval_expr(e));
            match &d.name {
                Pat::Ident(i) => {
                    let name = i.id.sym.to_string();
                    let rhs_alias = rhs.and_then(|r| self.node_binding_alias(r));
                    let v = if let Some(r) = rhs {
                        let id = self.graph.add_node(NodeKind::Variable, name.clone(), None);
                        self.flow(r, id);
                        if let Some(alias) = rhs_alias {
                            if let Some(node) = self.graph.node_mut(id) {
                                node.alias = Some(alias);
                            }
                        }
                        self.materialize_composite_bindings(&name, r);
                        id
                    } else {
                        self.graph.add_node(NodeKind::Variable, name.clone(), None)
                    };
                    self.cur().define(name.clone(), v);
                    if let Some(ref params) = self.fn_expr_params_at_depth_0.take() {
                        self.function_params.insert(name.clone(), params.clone());
                    }
                    if let Some(expr) = d.init.as_deref() {
                        if let Some(s) = self.try_infer_string_literal(expr) {
                            self.known_strings.insert(name.clone(), s);
                        }
                        if let Some(arr) = self.try_infer_array_literal(expr) {
                            self.known_arrays.insert(name.clone(), arr);
                        }
                    }
                }
                Pat::Array(_) | Pat::Object(_) => {
                    if let Some(r) = rhs {
                        self.bind_pat_with_source(&d.name, Some(r));
                    } else {
                        self.bind_pat_with_source(&d.name, None);
                    }
                }
                _ => {}
            }
        }
    }
}