use fxrank_core::effect::{RiskFeature, RiskKind, Tier};
use fxrank_core::score::weight_for_class;
use libcst_native::{Arg, Assert, AssignTargetExpression, Call, Expression, Raise};
use super::{
EffectSink,
expr::{leftmost_name, render_expr},
walk_own_body,
};
use crate::functions::FnUnit;
use crate::imports::Imports;
use crate::source::{SpanIndex, anchor_of_subslice};
pub fn detect(unit: &FnUnit, imports: &Imports, span: &SpanIndex, path: &str) -> Vec<RiskFeature> {
let mut sink = RiskSink {
imports,
span,
path: path.to_owned(),
features: Vec::new(),
};
walk_own_body(unit, &mut sink);
sink.features
}
struct RiskSink<'a> {
imports: &'a Imports,
span: &'a SpanIndex<'a>,
path: String,
features: Vec<RiskFeature>,
}
impl RiskSink<'_> {
fn push(&mut self, kind: RiskKind, tier: Tier, line: usize, evidence: String) {
let class = kind.class();
self.features.push(RiskFeature {
kind,
class,
weight: weight_for_class(class),
path: self.path.clone(),
line,
evidence,
tier,
});
}
fn resolve_dotted(&self, rendered: &str) -> Option<String> {
let (root, rest) = match rendered.split_once('.') {
Some((r, rest)) => (r, Some(rest)),
None => (rendered, None),
};
let base = self.imports.resolve(root)?;
Some(match rest {
Some(rest) => format!("{base}.{rest}"),
None => base.to_string(),
})
}
fn is_imported_name(&self, name: &str) -> bool {
self.imports.resolve(name).is_some()
}
}
impl EffectSink for RiskSink<'_> {
fn on_call(&mut self, call: &Call) {
let Some(rendered) = render_expr(&call.func) else {
return;
};
let line = leftmost_name(&call.func)
.map(|n| {
self.span
.line_col(anchor_of_subslice(self.span.src(), n.value))
.0
})
.unwrap_or(0);
match rendered.as_str() {
"eval" => {
self.push(
RiskKind::DynamicCode,
Tier::Exact,
line,
"eval(…) — dynamic code execution".into(),
);
return;
}
"exec" => {
self.push(
RiskKind::DynamicCode,
Tier::Exact,
line,
"exec(…) — dynamic code execution".into(),
);
return;
}
"compile" => {
self.push(
RiskKind::DynamicCode,
Tier::Exact,
line,
"compile(…) — dynamic code compilation".into(),
);
return;
}
"__import__" => {
self.push(
RiskKind::DynamicCode,
Tier::Exact,
line,
"__import__(…) — dynamic import".into(),
);
return;
}
_ => {}
}
if rendered == "setattr" {
if let Some(first_arg) = call.args.first() {
if let Expression::Name(n) = &first_arg.value
&& self.is_imported_name(n.value)
{
self.push(
RiskKind::DynamicCode,
Tier::Heuristic,
line,
format!("setattr({}, …) — monkey-patch on imported name", n.value),
);
}
}
return;
}
let resolved = self.resolve_dotted(&rendered);
if let Some(ref full) = resolved {
let root = full.split('.').next().unwrap_or(full.as_str());
if root == "subprocess" && has_shell_true(call) {
self.push(
RiskKind::DynamicCode,
Tier::Path,
line,
"subprocess(shell=True) — shell-injection surface".into(),
);
return;
}
}
if let Some(ref full) = resolved {
if matches!(full.as_str(), "pickle.load" | "pickle.loads") {
self.push(
RiskKind::DynamicCode,
Tier::Path,
line,
format!("{full}(…) — unsafe deserialization"),
);
return;
}
}
if let Some(ref full) = resolved {
if full == "yaml.load" {
self.push(
RiskKind::DynamicCode,
Tier::Path,
line,
"yaml.load(…) — unsafe YAML deserialization (use safe_load)".into(),
);
return;
}
}
if let Some(ref full) = resolved
&& full == "importlib.import_module"
{
self.push(
RiskKind::DynamicCode,
Tier::Path,
line,
"importlib.import_module(…) — dynamic import".into(),
);
}
}
fn on_assert(&mut self, _assert: &Assert) {}
fn on_raise(&mut self, _raise: &Raise) {}
fn on_assign_target(&mut self, _target: &AssignTargetExpression, _is_aug: bool) {}
}
fn has_shell_true(call: &Call) -> bool {
call.args.iter().any(|arg| is_shell_true_kwarg(arg))
}
fn is_shell_true_kwarg(arg: &Arg) -> bool {
let Some(kw) = &arg.keyword else { return false };
if kw.value != "shell" {
return false;
}
matches!(
&arg.value,
Expression::Name(n) if n.value == "True"
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::functions;
use crate::imports::Imports;
use crate::source::SpanIndex;
use std::collections::HashMap;
fn risk_features(name: &str) -> HashMap<String, Vec<String>> {
let src = std::fs::read_to_string(format!("tests/fixtures/{name}.py")).unwrap();
let module = libcst_native::parse_module(&src, None).unwrap();
let imports = Imports::build(&module);
let span = SpanIndex::new(&src);
let anchors = crate::source::lambda_anchors(&src).expect("tokenize must succeed");
let (units, _) = functions::collect(&module, &src, &span, &anchors);
let mut out: HashMap<String, Vec<String>> = HashMap::new();
for unit in &units {
let features = detect(unit, &imports, &span, "");
out.insert(
unit.symbol.clone(),
features.iter().map(|r| r.kind.wire().to_string()).collect(),
);
}
out
}
#[test]
fn detects_dynamic_code_and_shell() {
let r = risk_features("risk");
assert!(r["dyn"].contains(&"dynamic.code".to_string()));
assert!(r["deserialize"].contains(&"dynamic.code".to_string()));
assert!(r["shell"].contains(&"dynamic.code".to_string())); }
#[test]
fn detects_compile_and_dunder_import() {
let r = risk_features("risk");
assert!(
r["uses_compile"].contains(&"dynamic.code".to_string()),
"compile() must emit dynamic.code"
);
assert!(
r["uses_dunder_import"].contains(&"dynamic.code".to_string()),
"__import__() must emit dynamic.code"
);
}
#[test]
fn detects_yaml_load_but_not_safe_load() {
let r = risk_features("risk");
assert!(
r["unsafe_yaml"].contains(&"dynamic.code".to_string()),
"yaml.load() must emit dynamic.code"
);
assert!(
!r["safe_yaml"].contains(&"dynamic.code".to_string()),
"yaml.safe_load() must NOT emit dynamic.code"
);
}
#[test]
fn detects_importlib_import_module() {
let r = risk_features("risk");
assert!(
r["dynamic_import"].contains(&"dynamic.code".to_string()),
"importlib.import_module() must emit dynamic.code"
);
}
#[test]
fn detects_setattr_monkey_patch_on_imported_name_only() {
let r = risk_features("risk");
assert!(
r["monkey_patch"].contains(&"dynamic.code".to_string()),
"setattr on imported name must emit dynamic.code"
);
assert!(
!r["plain_setattr"].contains(&"dynamic.code".to_string()),
"setattr on non-imported name must NOT emit dynamic.code"
);
}
}