use std::sync::LazyLock;
use tree_sitter::{Node, Query, QueryCursor, StreamingIterator};
use crate::security::detect::powershell::ast::{CurrentAst, get_command_name, language};
use crate::security::detect::utils::{command_basename, node_extract_text};
use crate::security::detect::{Severity, ShellContext};
pub fn clean_ps_string(s: &str) -> String {
let s = s.trim();
if s.len() >= 2
&& ((s.starts_with('"') && s.ends_with('"')) || (s.starts_with('\'') && s.ends_with('\'')))
{
return s[1..s.len() - 1].to_string();
}
s.to_string()
}
pub fn extract_url_from_text(text: &str) -> Option<String> {
static URL_RE: LazyLock<regex::Regex> = LazyLock::new(|| {
regex::Regex::new("https?://[^\\s'\\\"`]+").expect("valid url regex")
});
URL_RE.find(text).map(|m| m.as_str().to_string())
}
pub fn is_ps_downloader(name: &str) -> bool {
let lower = name.to_lowercase();
let base = command_basename(&lower);
matches!(
base.as_str(),
"iwr"
| "invoke-webrequest"
| "curl"
| "wget"
| "curl.exe"
| "downloadstring"
| "downloaddata"
| "downloadfile"
)
}
static COMMON_QUERY: LazyLock<Query> = LazyLock::new(|| {
Query::new(
&language(),
"[(command) @cmd (redirection) @redirect (invokation_expression) @inv]",
)
.expect("invalid query")
});
pub async fn best_hit(
ctx: &ShellContext,
analyze: impl Fn(&Node, &[u8]) -> Option<(Severity, String)>,
) -> anyhow::Result<Option<(Severity, String)>> {
let current = ctx
.extensions
.get::<CurrentAst>()
.ok_or_else(|| anyhow::anyhow!("CurrentAst missing"))?;
let blocks = current.blocks.read().await;
let query = &*COMMON_QUERY;
let mut best: Option<(Severity, String)> = None;
for block in blocks.iter() {
let source_bytes = block.source.as_bytes();
let mut cursor = QueryCursor::new();
let mut matches = cursor.matches(query, block.tree.root_node(), source_bytes);
while let Some(m) = StreamingIterator::next(&mut matches) {
for capture in m.captures {
if let Some(hit) = analyze(&capture.node, source_bytes)
&& best.as_ref().is_none_or(|(b, _)| hit.0 > *b)
{
best = Some(hit);
}
}
}
}
Ok(best)
}
pub async fn best_text(
ctx: &ShellContext,
analyze: impl Fn(&str) -> Option<(Severity, String)>,
) -> anyhow::Result<Option<(Severity, String)>> {
let current = ctx
.extensions
.get::<CurrentAst>()
.ok_or_else(|| anyhow::anyhow!("CurrentAst missing"))?;
let blocks = current.blocks.read().await;
let mut best: Option<(Severity, String)> = None;
for block in blocks.iter() {
if let Some(hit) = analyze(&block.source)
&& best.as_ref().is_none_or(|(b, _)| hit.0 > *b)
{
best = Some(hit);
}
}
Ok(best)
}
pub fn collect_args<'a>(node: &Node, source: &'a [u8]) -> Vec<&'a str> {
collect_elements(node, source)
.into_iter()
.filter(|(kind, _)| !matches!(kind.as_str(), "command_parameter" | "redirection"))
.map(|(_, text)| text)
.collect()
}
pub fn collect_elements<'a>(node: &Node, source: &'a [u8]) -> Vec<(String, &'a str)> {
let mut out = Vec::new();
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
if child.kind() == "command_elements" {
let mut ec = child.walk();
for el in child.children(&mut ec) {
if el.kind() == "command_argument_sep" {
continue;
}
if let Some(t) = node_extract_text(&el, source) {
out.push((el.kind().to_string(), t));
}
}
}
}
out
}
pub fn collect_parameters<'a>(node: &Node, source: &'a [u8]) -> Vec<&'a str> {
collect_elements(node, source)
.into_iter()
.filter(|(kind, _)| kind == "command_parameter")
.map(|(_, text)| text)
.collect()
}
pub fn redirect_target<'a>(node: &Node, source: &'a [u8]) -> Option<(&'a str, &'a str)> {
if node.kind() != "redirection" {
return None;
}
let mut operator: Option<&str> = None;
let mut dest: Option<&'a str> = None;
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
match child.kind() {
"file_redirection_operator" => {
operator = node_extract_text(&child, source);
}
"redirected_file_name" => {
dest = node_extract_text(&child, source);
}
_ => {}
}
}
match (operator, dest) {
(Some(op), Some(dest)) => Some((op, dest)),
_ => None,
}
}
pub fn ps_normalize_path(t: &str) -> String {
let mut s = t.trim().trim_matches(|c| c == '\'' || c == '"').to_string();
if s.len() >= 2 && s.as_bytes()[1] == b':' {
let drive = s.as_bytes()[0].to_ascii_lowercase() as char;
s.replace_range(0..1, &drive.to_string());
}
while s.len() > 1 && (s.ends_with('\\') || s.ends_with('/')) {
s.pop();
}
s
}
pub fn ps_path_has_marker(p: &str, markers: &[&str]) -> bool {
let s = ps_normalize_path(p).to_lowercase();
markers.iter().any(|m| s.contains(&m.to_lowercase()))
}
pub fn has_parameter(args: &[&str], param: &str) -> bool {
let param_lower = param.to_lowercase();
args.iter().any(|a| {
let t = a.trim_start_matches('-').trim_start_matches('/');
t.to_lowercase() == param_lower
})
}
pub fn has_any_parameter(args: &[&str], params: &[&str]) -> bool {
params.iter().any(|p| has_parameter(args, p))
}
pub fn is_ps_shell_sink(name: &str) -> bool {
let lower = name.to_lowercase();
matches!(
lower.as_str(),
"iex"
| "invoke-expression"
| "powershell"
| "pwsh"
| "powershell.exe"
| "pwsh.exe"
| "cmd"
| "cmd.exe"
)
}
pub fn detect_remote_execution_text(text: &str) -> Option<(String, String)> {
let lower = text.to_lowercase();
let sink = if lower.contains("invoke-expression") || lower.contains("iex") {
"iex"
} else if lower.contains("pwsh") || lower.contains("powershell") {
"powershell"
} else if lower.contains("cmd") {
"cmd"
} else {
return None;
};
let url = extract_url_from_text(text)?;
Some((sink.to_string(), url))
}
pub fn extract_downloader_url(node: &Node, source: &[u8]) -> Option<String> {
let mut stack = vec![*node];
while let Some(n) = stack.pop() {
if n.kind() == "command"
&& let Some(cmd_name) = get_command_name(&n, source)
&& is_ps_downloader(cmd_name)
{
for arg in collect_args(&n, source) {
let cleaned = clean_ps_string(arg);
if cleaned.starts_with("http://") || cleaned.starts_with("https://") {
return Some(cleaned);
}
}
}
let mut cursor = n.walk();
for child in n.children(&mut cursor) {
stack.push(child);
}
}
None
}