Skip to content

Commit

Permalink
update .Net features to capa 6.1.0
Browse files Browse the repository at this point in the history
  • Loading branch information
@mnaza
mnaza committed Oct 25, 2023
1 parent 1c8c3c7 commit 3a4761e
Show file tree
Hide file tree
Showing 6 changed files with 190 additions and 37 deletions.
4 changes: 2 additions & 2 deletions Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ pdb = "0.8.0"
petgraph = "0.6.2"
regex = "1.5"
serde = { version = "1", features = ["derive"] }
smda = "0.2.5"
smda = { git = "https://github.com/marirs/smda-rs.git" }
thiserror = "1"
walkdir = "2.3.2"
yaml-rust = "0.4.5"
Expand All @@ -30,7 +30,7 @@ lazy_static = "1.4.0"
[dev-dependencies]
clap = { version = "4.0.27", features = ["cargo", "derive"] }
serde_json = "1"
prettytable-rs = "0.9.0"
prettytable-rs = "0.10.0"

[lib]
name = "capa"
Expand Down
2 changes: 1 addition & 1 deletion src/consts.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use serde::{Deserialize, Serialize};
use std::fmt::Display;

#[derive(Debug, Clone, Serialize, Deserialize, Copy)]
#[derive(Debug, Clone, Serialize, Deserialize, Copy, PartialEq, Eq)]
#[allow(clippy::upper_case_acronyms)]
pub enum FileFormat {
PE,
Expand Down
166 changes: 153 additions & 13 deletions src/extractor/dnfile.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,10 @@ use dnfile::{
stream::meta_data_tables::mdtables::{codedindex::*, *},
DnPe,
};
use std::collections::HashMap;
use std::{
collections::{HashMap, HashSet},
sync::{Arc, RwLock},
};

use crate::extractor::Extractor as BaseExtractor;

Expand All @@ -31,6 +34,8 @@ impl super::Instruction for Instruction {
#[derive(Debug, Clone)]
struct Function {
f: cil::function::Function,
calls_to: HashSet<u64>,
calls_from: HashSet<u64>,
}

impl super::Function for Function {
Expand Down Expand Up @@ -102,6 +107,8 @@ enum Callee {

#[derive(Debug)]
pub struct Extractor {
properties_cache: Arc<RwLock<Option<HashMap<u64, DnMethod>>>>,
fields_cache: Arc<RwLock<Option<HashMap<u64, DnMethod>>>>,
pe: DnPe,
}

Expand Down Expand Up @@ -147,27 +154,69 @@ impl super::Extractor for Extractor {
fn extract_file_features(&self) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
let mut ss = self.extract_file_import_names()?;
ss.extend(self.extract_file_function_names()?);
//ss.extend(self.extract_file_string()?);
ss.extend(self.extract_file_format()?);
ss.extend(self.extract_file_mixed_mode_characteristic_features()?);
ss.extend(self.extract_file_namespace_and_class_features()?);
Ok(ss)
}

fn get_functions(&self) -> Result<std::collections::HashMap<u64, Box<dyn super::Function>>> {
let mut res: std::collections::HashMap<u64, Box<dyn super::Function>> =
let mut methods: std::collections::HashMap<u64, Function> =
std::collections::HashMap::new();
let mut calls_to_map = HashMap::new();
for f in self.pe.net()?.functions() {
res.insert(f.offset as u64, Box::new(Function { f: f.clone() }));
let mut calls_from = HashSet::new();
for insn in &f.instructions {
if ![
OpCodeValue::Call,
OpCodeValue::Callvirt,
OpCodeValue::Jmp,
OpCodeValue::Newobj,
]
.contains(&insn.opcode.value)
{
continue;
}
let address = insn.operand.value()?;
let ee = calls_to_map.entry(address as u64).or_insert(HashSet::new());
ee.insert(f.offset as u64);
calls_from.insert(address as u64);
}
methods.insert(
f.offset as u64,
Function {
f: f.clone(),
calls_to: HashSet::new(),
calls_from,
},
);
}
Ok(res)
for (a, calls_from) in calls_to_map.into_iter() {
if let Some(f) = methods.get_mut(&a) {
f.calls_from = calls_from;
}
}
Ok(methods
.into_iter()
.map(|(a, b)| (a, Box::new(b) as Box<dyn super::Function>))
.collect::<HashMap<u64, Box<dyn super::Function>>>())
}

fn extract_function_features(
&self,
_f: &Box<dyn super::Function>,
f: &Box<dyn super::Function>,
) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
Ok(vec![])
let f: &Function = f.as_any().downcast_ref::<Function>().unwrap();
Ok([
self.extract_function_call_to_features(&f)?,
self.extract_function_call_from_features(&f)?,
self.extract_recurcive_call_features(&f)?,
]
.into_iter()
.fold(Vec::new(), |mut acc, f| {
acc.extend(f);
acc
}))
}

fn get_basic_blocks(
Expand Down Expand Up @@ -215,6 +264,8 @@ impl Extractor {
pub fn new(file_path: &str) -> Result<Extractor> {
let res = Extractor {
pe: dnfile::DnPe::new(file_path)?,
fields_cache: Arc::new(RwLock::new(None)),
properties_cache: Arc::new(RwLock::new(None)),
};
Ok(res)
}
Expand Down Expand Up @@ -381,6 +432,13 @@ impl Extractor {
Ok(None)
}

pub fn get_properties(&self) -> Result<&RwLock<Option<HashMap<u64, DnMethod>>>> {
if let None = *self.properties_cache.read().unwrap() {
*self.properties_cache.write().unwrap() = Some(self.get_dotnet_properties()?);
}
Ok(&self.properties_cache)
}

pub fn get_dotnet_properties(&self) -> Result<HashMap<u64, DnMethod>> {
let mut res = HashMap::new();
let method_semantics = if let Ok(s) = self.pe.net()?.md_table("MethodSemantics") {
Expand Down Expand Up @@ -430,6 +488,13 @@ impl Extractor {
Ok(res)
}

pub fn get_fields(&self) -> Result<&RwLock<Option<HashMap<u64, DnMethod>>>> {
if let None = *self.fields_cache.read().unwrap() {
*self.fields_cache.write().unwrap() = Some(self.get_dotnet_fields()?);
}
Ok(&self.fields_cache)
}

/// get fields from TypeDef table
pub fn get_dotnet_fields(&self) -> Result<HashMap<u64, DnMethod>> {
let mut res = HashMap::new();
Expand Down Expand Up @@ -507,6 +572,57 @@ impl Extractor {
}
}

fn extract_function_call_to_features(
&self,
f: &Function,
) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
Ok(f.calls_to
.iter()
.map(|a| {
(
crate::rules::features::Feature::Characteristic(
crate::rules::features::CharacteristicFeature::new("calls to", "").unwrap(),
),
*a,
)
})
.collect())
}

fn extract_function_call_from_features(
&self,
f: &Function,
) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
Ok(f.calls_to
.iter()
.map(|a| {
(
crate::rules::features::Feature::Characteristic(
crate::rules::features::CharacteristicFeature::new("calls from", "")
.unwrap(),
),
*a,
)
})
.collect())
}

fn extract_recurcive_call_features(
&self,
f: &Function,
) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
if f.calls_to.contains(&(f.f.offset as u64)) {
Ok(vec![(
crate::rules::features::Feature::Characteristic(
crate::rules::features::CharacteristicFeature::new("recursive call", "")?,
),
f.f.offset as u64,
)])
} else {
Ok(vec![])
}
}

fn get_callee(&self, token: u64) -> Result<Option<Callee>> {
// map dotnet token to un/managed method
match self.get_dotnet_managed_imports()?.get(&token) {
Expand Down Expand Up @@ -569,7 +685,11 @@ impl Extractor {
)?;
if row.downcast_ref::<MethodDef>().is_some() {
if self
.get_dotnet_properties()?
.get_properties()?
.read()
.unwrap()
.as_ref()
.unwrap()
.get(&(insn.operand.value()? as u64))
.is_some()
{
Expand Down Expand Up @@ -611,7 +731,11 @@ impl Extractor {
if let Some(_operand) = operand.downcast_ref::<MethodDef>() {
// check if the method belongs to the MethodDef table and whether it is used to access a property
if let Some(prop) = self
.get_dotnet_properties()?
.get_properties()?
.read()
.unwrap()
.as_ref()
.unwrap()
.get(&(insn.operand.value()? as u64))
{
return Ok(vec![(
Expand Down Expand Up @@ -688,7 +812,11 @@ impl Extractor {
if let Some(_operand) = operand.downcast_ref::<Field>() {
// determine whether the operand is a field by checking if it belongs to the Field table
if let Some(field) = self
.get_dotnet_fields()?
.get_fields()?
.read()
.unwrap()
.as_ref()
.unwrap()
.get(&(insn.operand.value()? as u64))
{
return Ok(vec![(
Expand All @@ -711,7 +839,11 @@ impl Extractor {
if let Some(_operand) = operand.downcast_ref::<Field>() {
// determine whether the operand is a field by checking if it belongs to the Field table
if let Some(field) = self
.get_dotnet_fields()?
.get_fields()?
.read()
.unwrap()
.as_ref()
.unwrap()
.get(&(insn.operand.value()? as u64))
{
return Ok(vec![(
Expand Down Expand Up @@ -833,7 +965,11 @@ impl Extractor {
}
} else if operand.downcast_ref::<Field>().is_some() {
if let Some(field) = self
.get_dotnet_fields()?
.get_fields()?
.read()
.unwrap()
.as_ref()
.unwrap()
.get(&(insn.operand.value()? as u64))
{
res.push((
Expand Down Expand Up @@ -911,7 +1047,11 @@ impl Extractor {
}
} else if operand.downcast_ref::<Field>().is_some() {
if let Some(field) = self
.get_dotnet_fields()?
.get_fields()?
.read()
.unwrap()
.as_ref()
.unwrap()
.get(&(insn.operand.value()? as u64))
{
res.push((
Expand Down
53 changes: 32 additions & 21 deletions src/extractor/smda.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ use smda::{
};
use std::{collections::HashMap, convert::TryInto};

//use super::Extractor;
use crate::extractor::Extractor as _;

#[derive(Debug, Clone)]
struct InstructionS {
i: Instruction,
Expand Down Expand Up @@ -131,8 +134,7 @@ impl super::Extractor for Extractor {
&self,
f: &Box<dyn super::Function>,
) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
let mut res = vec![];
//extract function calls to
let mut res = vec![]; //extract function calls to
for inref in f.inrefs() {
res.push((
crate::rules::features::Feature::Characteristic(
Expand Down Expand Up @@ -310,18 +312,6 @@ impl Extractor {
Ok(res)
}

// pub fn extract_file_features(&self) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
// let mut res = vec![];
// // res.extend(self.extract_file_embedded_pe()?);
// res.extend(self.extract_file_export_names()?);
// res.extend(self.extract_file_import_names()?);
// res.extend(self.extract_file_section_names()?);
// res.extend(self.extract_file_strings()?);
// // res.extend(self.extract_file_function_names(pbytes)?);
// res.extend(self.extract_file_format()?);
// Ok(res)
// }

fn extract_file_format(&self) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
let mut res = vec![];
res.push((
Expand Down Expand Up @@ -368,13 +358,34 @@ impl Extractor {

pub fn extract_file_export_names(&self) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
let mut res = vec![];
for (e, o) in &self.report.exports {
res.push((
crate::rules::features::Feature::Export(
crate::rules::features::ExportFeature::new(e, "")?,
),
*o as u64,
));
for (e, o, ree) in &self.report.exports {
match ree {
None => {
res.push((
crate::rules::features::Feature::Export(
crate::rules::features::ExportFeature::new(e, "")?,
),
*o as u64 - self.get_base_address()?,
));
}
Some(re) => {
res.push((
crate::rules::features::Feature::Export(
crate::rules::features::ExportFeature::new(re, "")?,
),
*o as u64 - self.get_base_address()?,
));
res.push((
crate::rules::features::Feature::Characteristic(
crate::rules::features::CharacteristicFeature::new(
"forwarded export",
"",
)?,
),
*o as u64 - self.get_base_address()?,
));
}
}
}
Ok(res)
}
Expand Down
1 change: 1 addition & 0 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -194,6 +194,7 @@ impl FileCapabilities {
if addr == &0 {
continue;
}
eprintln!("{}, {}", arrd, count);
let mut fc = FunctionCapabilities {
address: *addr as usize,
features: *count,
Expand Down
Loading

0 comments on commit 3a4761e

Please sign in to comment.