use std::cell::RefCell;
use std::rc::Rc;
use std::slice::Iter;
use bstr::BStr;
use digest::Digest;
use itertools::Itertools;
use nom::branch::alt;
use nom::character::complete::u8;
use nom::combinator::map;
use nom::number::complete::{le_u16, le_u32};
use crate::compiler::RegexpId;
use crate::modules::prelude::*;
use crate::modules::protos::pe::*;
use crate::types::Struct;
#[cfg(test)]
mod tests;
mod authenticode;
pub mod parser;
mod rva2off;
thread_local!(
static IMPHASH_CACHE: RefCell<Option<String>> =
const { RefCell::new(None) };
static CHECKSUM_CACHE: RefCell<Option<i64>> = const { RefCell::new(None) };
);
#[module_main]
fn main(data: &[u8], _meta: Option<&[u8]>) -> Result<PE, ModuleError> {
IMPHASH_CACHE.with(|cache| *cache.borrow_mut() = None);
CHECKSUM_CACHE.with(|cache| *cache.borrow_mut() = None);
match parser::PE::parse(data) {
Ok(pe) => Ok(pe.into()),
Err(_) => {
let mut pe = PE::new();
pe.is_pe = Some(false);
Ok(pe)
}
}
}
#[module_export]
fn is_32bit(ctx: &ScanContext) -> Option<bool> {
let magic = ctx.module_output::<PE>()?.opthdr_magic?;
Some(magic.value() == OptionalMagic::IMAGE_NT_OPTIONAL_HDR32_MAGIC as i32)
}
#[module_export]
fn is_64bit(ctx: &ScanContext) -> Option<bool> {
let magic = ctx.module_output::<PE>()?.opthdr_magic?;
Some(magic.value() == OptionalMagic::IMAGE_NT_OPTIONAL_HDR64_MAGIC as i32)
}
#[module_export]
fn is_dll(ctx: &ScanContext) -> Option<bool> {
let characteristics = ctx.module_output::<PE>()?.characteristics?;
Some(characteristics & Characteristics::DLL as u32 != 0)
}
#[module_export]
fn rva_to_offset(ctx: &ScanContext, rva: i64) -> Option<i64> {
let pe = ctx.module_output::<PE>()?;
let offset = rva2off::rva_to_offset(
rva.try_into().ok()?,
pe.sections.as_slice(),
pe.file_alignment?,
pe.section_alignment?,
)?;
Some(offset.into())
}
#[module_export]
fn calculate_checksum(ctx: &mut ScanContext) -> Option<i64> {
let cached: Option<i64> =
CHECKSUM_CACHE.with(|cache| -> Option<i64> { *cache.borrow() });
if cached.is_some() {
return cached;
}
let pe = ctx.module_output::<PE>()?;
let data = ctx.scanned_data()?;
let mut sum: u32 = 0;
if !pe.is_pe() {
return None;
}
let data_parser = alt((
le_u32::<&[u8], nom::error::Error<&[u8]>>,
map(le_u16, |v| v as u32),
map(u8, |v| v as u32),
));
for v in &mut nom::combinator::iterator(data, data_parser) {
sum = match sum.overflowing_add(v) {
(s, true) => s + 1, (s, false) => s, }
}
sum = match sum.overflowing_sub(pe.checksum?) {
(s, true) => s - 1, (s, false) => s, };
sum = (sum & 0xffff) + (sum >> 16);
sum += sum >> 16;
sum &= 0xffff;
sum += data.len() as u32;
CHECKSUM_CACHE.with(|cache| {
*cache.borrow_mut() = Some(sum.into());
});
Some(sum.into())
}
#[module_export(name = "section_index")]
fn section_index_name(ctx: &ScanContext, name: RuntimeString) -> Option<i64> {
let pe = ctx.module_output::<PE>()?;
let name = name.as_bstr(ctx);
pe.sections
.iter()
.find_position(|section| {
section.name.as_deref().is_some_and(|n| n == name.as_bytes())
})
.map(|(index, _)| index as i64)
}
#[module_export(name = "section_index")]
fn section_index_offset(ctx: &ScanContext, offset: i64) -> Option<i64> {
let pe = ctx.module_output::<PE>()?;
let offset: u32 = offset.try_into().ok()?;
pe.sections
.iter()
.find_position(|section| {
match (section.raw_data_offset, section.raw_data_size) {
(Some(section_offset), Some(section_size)) => (section_offset
..section_offset + section_size)
.contains(&offset),
_ => false,
}
})
.map(|(index, _)| index as i64)
}
#[module_export]
fn imphash(ctx: &mut ScanContext) -> Option<Lowercase<FixedLenString<32>>> {
let cached =
IMPHASH_CACHE.with(|cache| -> Option<Lowercase<FixedLenString<32>>> {
cache.borrow().as_deref().map(|s| {
Lowercase::<FixedLenString<32>>::from_slice(ctx, s.as_bytes())
})
});
if cached.is_some() {
return cached;
}
let pe = ctx.module_output::<PE>()?;
if !pe.is_pe() {
return None;
}
let mut md5_hash = md5::Md5::default();
let mut first = true;
for import in &pe.import_details {
let original_dll_name =
import.library_name.as_deref().unwrap().to_lowercase();
let mut dll_name = original_dll_name.as_str();
for extension in [".dll", ".sys", ".ocx"] {
dll_name = dll_name.trim_end_matches(extension);
}
for func in &import.functions {
if !first {
Digest::update(&mut md5_hash, ",".as_bytes())
}
Digest::update(&mut md5_hash, dll_name);
Digest::update(&mut md5_hash, ".".as_bytes());
Digest::update(
&mut md5_hash,
func.name.as_deref().unwrap().to_lowercase().as_bytes(),
);
first = false;
}
}
let digest = format!("{:x}", md5_hash.finalize());
IMPHASH_CACHE.with(|cache| {
*cache.borrow_mut() = Some(digest.clone());
});
Some(Lowercase::<FixedLenString<32>>::new(digest))
}
#[module_export(name = "rich_signature.toolid")]
fn rich_toolid(ctx: &mut ScanContext, toolid: i64) -> Option<i64> {
rich_version_impl(ctx.module_output::<PE>()?, Some(toolid), None)
}
#[module_export(name = "rich_signature.version")]
fn rich_version(ctx: &mut ScanContext, version: i64) -> Option<i64> {
rich_version_impl(ctx.module_output::<PE>()?, None, Some(version))
}
#[module_export(name = "rich_signature.version")]
fn rich_version_toolid(
ctx: &mut ScanContext,
version: i64,
toolid: i64,
) -> Option<i64> {
rich_version_impl(ctx.module_output::<PE>()?, Some(toolid), Some(version))
}
#[module_export(name = "rich_signature.toolid")]
fn rich_toolid_version(
ctx: &mut ScanContext,
toolid: i64,
version: i64,
) -> Option<i64> {
rich_version_impl(ctx.module_output::<PE>()?, Some(toolid), Some(version))
}
fn rich_version_impl(
pe: &PE,
toolid: Option<i64>,
version: Option<i64>,
) -> Option<i64> {
assert!(toolid.is_some() || version.is_some());
let count = pe
.rich_signature
.tools
.iter()
.filter_map(|t| {
let toolid_matches = toolid
.map(|toolid| toolid == t.toolid.unwrap() as i64)
.unwrap_or(true);
let version_matches = version
.map(|version| version == t.version.unwrap() as i64)
.unwrap_or(true);
if toolid_matches && version_matches {
t.times.map(|v| v as i64)
} else {
None
}
})
.sum::<i64>();
Some(count)
}
#[module_export(name = "imports")]
fn standard_imports_dll(
ctx: &ScanContext,
dll_name: RuntimeString,
) -> Option<i64> {
imports_impl(
ctx,
ImportFlags::IMPORT_STANDARD as i64,
MatchCriteria::Name(dll_name.as_bstr(ctx)),
MatchCriteria::Any,
)
}
#[module_export(name = "imports")]
fn standard_imports_func(
ctx: &ScanContext,
dll_name: RuntimeString,
func_name: RuntimeString,
) -> Option<bool> {
Some(
imports_impl(
ctx,
ImportFlags::IMPORT_STANDARD as i64,
MatchCriteria::Name(dll_name.as_bstr(ctx)),
MatchCriteria::Name(func_name.as_bstr(ctx)),
)? > 0,
)
}
#[module_export(name = "imports")]
fn standard_imports_ordinal(
ctx: &ScanContext,
dll_name: RuntimeString,
ordinal: i64,
) -> Option<i64> {
imports_impl(
ctx,
ImportFlags::IMPORT_STANDARD as i64,
MatchCriteria::Name(dll_name.as_bstr(ctx)),
MatchCriteria::Ordinal(ordinal),
)
}
#[module_export(name = "imports")]
fn standard_imports_regexp(
ctx: &ScanContext,
dll_name: RegexpId,
func_name: RegexpId,
) -> Option<i64> {
imports_impl(
ctx,
ImportFlags::IMPORT_STANDARD as i64,
MatchCriteria::Regexp(dll_name),
MatchCriteria::Regexp(func_name),
)
}
#[module_export(name = "imports")]
fn imports_dll(
ctx: &ScanContext,
import_flags: i64,
dll_name: RuntimeString,
) -> Option<i64> {
imports_impl(
ctx,
import_flags,
MatchCriteria::Name(dll_name.as_bstr(ctx)),
MatchCriteria::Any,
)
}
#[module_export(name = "imports")]
fn imports_func(
ctx: &ScanContext,
import_flags: i64,
dll_name: RuntimeString,
func_name: RuntimeString,
) -> Option<bool> {
Some(
imports_impl(
ctx,
import_flags,
MatchCriteria::Name(dll_name.as_bstr(ctx)),
MatchCriteria::Name(func_name.as_bstr(ctx)),
)? > 0,
)
}
#[module_export(name = "imports")]
fn imports_ordinal(
ctx: &ScanContext,
import_flags: i64,
dll_name: RuntimeString,
ordinal: i64,
) -> Option<bool> {
Some(
imports_impl(
ctx,
import_flags,
MatchCriteria::Name(dll_name.as_bstr(ctx)),
MatchCriteria::Ordinal(ordinal),
)? > 0,
)
}
#[module_export(name = "imports")]
fn imports_regexp(
ctx: &ScanContext,
import_flags: i64,
dll_name: RegexpId,
func_name: RegexpId,
) -> Option<i64> {
imports_impl(
ctx,
import_flags,
MatchCriteria::Regexp(dll_name),
MatchCriteria::Regexp(func_name),
)
}
#[module_export(name = "import_rva")]
fn import_rva_func(
ctx: &ScanContext,
dll_name: RuntimeString,
func_name: RuntimeString,
) -> Option<i64> {
let pe = ctx.module_output::<PE>()?;
import_rva_impl(
pe.import_details.as_slice(),
MatchCriteria::Name(dll_name.as_bstr(ctx)),
MatchCriteria::Name(func_name.as_bstr(ctx)),
)
}
#[module_export(name = "import_rva")]
fn import_rva_ordinal(
ctx: &ScanContext,
dll_name: RuntimeString,
ordinal: i64,
) -> Option<i64> {
let pe = ctx.module_output::<PE>()?;
import_rva_impl(
pe.import_details.as_slice(),
MatchCriteria::Name(dll_name.as_bstr(ctx)),
MatchCriteria::Ordinal(ordinal),
)
}
#[module_export(name = "delayed_import_rva")]
fn delayed_import_rva_func(
ctx: &ScanContext,
dll_name: RuntimeString,
func_name: RuntimeString,
) -> Option<i64> {
let pe = ctx.module_output::<PE>()?;
import_rva_impl(
pe.delayed_import_details.as_slice(),
MatchCriteria::Name(dll_name.as_bstr(ctx)),
MatchCriteria::Name(func_name.as_bstr(ctx)),
)
}
#[module_export(name = "delayed_import_rva")]
fn delayed_import_rva_ordinal(
ctx: &ScanContext,
dll_name: RuntimeString,
ordinal: i64,
) -> Option<i64> {
let pe = ctx.module_output::<PE>()?;
import_rva_impl(
pe.delayed_import_details.as_slice(),
MatchCriteria::Name(dll_name.as_bstr(ctx)),
MatchCriteria::Ordinal(ordinal),
)
}
#[module_export(name = "exports")]
fn exports_func(ctx: &ScanContext, func_name: RuntimeString) -> Option<bool> {
let (found, _) =
exports_impl(ctx, MatchCriteria::Name(func_name.as_bstr(ctx)))?;
Some(found)
}
#[module_export(name = "exports")]
fn exports_ordinal(ctx: &ScanContext, ordinal: i64) -> Option<bool> {
let (found, _) = exports_impl(ctx, MatchCriteria::Ordinal(ordinal))?;
Some(found)
}
#[module_export(name = "exports")]
fn exports_regexp(ctx: &ScanContext, func_name: RegexpId) -> Option<bool> {
let (found, _) = exports_impl(ctx, MatchCriteria::Regexp(func_name))?;
Some(found)
}
#[module_export(name = "exports_index")]
fn exports_index_func(
ctx: &ScanContext,
func_name: RuntimeString,
) -> Option<i64> {
match exports_impl(ctx, MatchCriteria::Name(func_name.as_bstr(ctx))) {
Some((true, position)) => Some(position as i64),
_ => None,
}
}
#[module_export(name = "exports_index")]
fn exports_index_ordinal(ctx: &ScanContext, ordinal: i64) -> Option<i64> {
match exports_impl(ctx, MatchCriteria::Ordinal(ordinal)) {
Some((true, position)) => Some(position as i64),
_ => None,
}
}
#[module_export(name = "exports_index")]
fn exports_index_regexp(
ctx: &ScanContext,
func_name: RegexpId,
) -> Option<i64> {
match exports_impl(ctx, MatchCriteria::Regexp(func_name)) {
Some((true, position)) => Some(position as i64),
_ => None,
}
}
#[module_export]
fn locale(ctx: &ScanContext, loc: i64) -> Option<bool> {
let pe = ctx.module_output::<PE>()?;
let loc: u32 = match loc.try_into() {
Ok(lang) => lang,
Err(_) => return Some(false),
};
Some(pe.resources.iter().any(|resource| {
resource.language.is_some_and(|rsrc_lang| rsrc_lang & 0xffff == loc)
}))
}
#[module_export]
fn language(ctx: &ScanContext, lang: i64) -> Option<bool> {
let pe = ctx.module_output::<PE>()?;
let lang: u32 = match lang.try_into() {
Ok(lang) => lang,
Err(_) => return Some(false),
};
Some(pe.resources.iter().any(|resource| {
resource.language.is_some_and(|rsrc_lang| rsrc_lang & 0xff == lang)
}))
}
#[module_export(method_of = "pe.Signature")]
fn valid_on(
_ctx: &ScanContext,
signature: Rc<Struct>,
timestamp: i64,
) -> Option<bool> {
let not_before = signature
.field_by_name("not_before")
.unwrap()
.type_value
.try_as_integer()?;
let not_after = signature
.field_by_name("not_after")
.unwrap()
.type_value
.try_as_integer()?;
Some(timestamp >= not_before && timestamp <= not_after)
}
enum MatchCriteria<'a> {
Any,
Regexp(RegexpId),
Name(&'a BStr),
Ordinal(i64),
}
fn imports_impl(
ctx: &ScanContext,
import_flags: i64,
expected_dll_name: MatchCriteria,
expected_func_name: MatchCriteria,
) -> Option<i64> {
let count_matching_funcs = |it: Iter<'_, Function>| {
it.filter(|func| match expected_func_name {
MatchCriteria::Any => true,
MatchCriteria::Name(expected_name) => {
func.name.as_ref().is_some_and(|name| {
expected_name.eq_ignore_ascii_case(name.as_bytes())
})
}
MatchCriteria::Regexp(regexp_id) => {
func.name.as_ref().is_some_and(|name| {
ctx.regexp_matches(regexp_id, name.as_bytes())
})
}
MatchCriteria::Ordinal(expected_ordinal) => func
.ordinal
.is_some_and(|ordinal| ordinal as i64 == expected_ordinal),
})
.count()
};
let count_matching_imports = |it: Iter<'_, Import>| {
it.filter_map(|import| {
let name_matches = match expected_dll_name {
MatchCriteria::Any => true,
MatchCriteria::Name(expected_name) => {
import.library_name.as_ref().is_some_and(|name| {
expected_name.eq_ignore_ascii_case(name.as_bytes())
})
}
MatchCriteria::Regexp(regexp_id) => {
import.library_name.as_ref().is_some_and(|name| {
ctx.regexp_matches(regexp_id, name.as_bytes())
})
}
MatchCriteria::Ordinal(_) => unreachable!(),
};
if name_matches {
Some(count_matching_funcs(import.functions.iter()))
} else {
None
}
})
.sum::<usize>()
};
let pe = ctx.module_output::<PE>()?;
let mut total = 0;
if import_flags & ImportFlags::IMPORT_STANDARD as i64 != 0 {
total += count_matching_imports(pe.import_details.iter());
}
if import_flags & ImportFlags::IMPORT_DELAYED as i64 != 0 {
total += count_matching_imports(pe.delayed_import_details.iter());
}
total.try_into().ok()
}
fn import_rva_impl(
imports: &[Import],
expected_dll_name: MatchCriteria,
expected_func_name: MatchCriteria,
) -> Option<i64> {
for import in imports {
let matches = match expected_dll_name {
MatchCriteria::Any => true,
MatchCriteria::Name(expected_name) => {
import.library_name.as_ref().is_some_and(|name| {
expected_name.eq_ignore_ascii_case(name.as_bytes())
})
}
MatchCriteria::Regexp(_) => unreachable!(),
MatchCriteria::Ordinal(_) => unreachable!(),
};
if matches {
for func in import.functions.iter() {
match expected_func_name {
MatchCriteria::Any => return func.rva.map(|r| r as i64),
MatchCriteria::Name(expected_name) => {
if func.name.as_ref().is_some_and(|name| {
expected_name.eq_ignore_ascii_case(name.as_bytes())
}) {
return func.rva.map(|r| r as i64);
}
}
MatchCriteria::Ordinal(expected_ordinal) => {
if func.ordinal.is_some_and(|ordinal| {
ordinal as i64 == expected_ordinal
}) {
return func.rva.map(|r| r as i64);
}
}
MatchCriteria::Regexp(_) => unreachable!(),
}
}
}
}
None
}
fn exports_impl(
ctx: &ScanContext,
expected_func_name: MatchCriteria,
) -> Option<(bool, usize)> {
let pe = ctx.module_output::<PE>()?;
pe.export_details
.iter()
.find_position(|export| match expected_func_name {
MatchCriteria::Any => true,
MatchCriteria::Regexp(regexp_id) => {
export.name.as_ref().is_some_and(|name| {
ctx.regexp_matches(regexp_id, name.as_bytes())
})
}
MatchCriteria::Name(expected_name) => {
export.name.as_ref().is_some_and(|name| {
expected_name.eq_ignore_ascii_case(name.as_bytes())
})
}
MatchCriteria::Ordinal(expected_ordinal) => export
.ordinal
.as_ref()
.is_some_and(|ordinal| expected_ordinal == *ordinal as i64),
})
.map_or(Some((false, 0)), |(position, _)| Some((true, position)))
}