Skip to content

Commit

Permalink
refactors
Browse files Browse the repository at this point in the history
  • Loading branch information
@marirs
marirs committed Feb 11, 2024
1 parent 47901db commit 8fd4803
Show file tree
Hide file tree
Showing 7 changed files with 211 additions and 156 deletions.
9 changes: 5 additions & 4 deletions Cargo.toml
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
[package]
name = "capa"
version = "0.3.8"
version = "0.3.9"
description = "File capability extractor."
authors = ["Marirs <[email protected]>", "Andrey Mnatsakanov <[email protected]>"]
keywords = ["capa", "fce", "capability", "file"]
Expand All @@ -19,14 +19,15 @@ petgraph = "0.6.2"
regex = "1.5"
fancy-regex = { git = "https://github.com/mnaza/fancy-regex.git" }
serde = { version = "1", features = ["derive"] }
smda = {git = "https://github.com/jorgeaduran/smda-rs.git", branch = "fixes"}
#smda = {git = "https://github.com/jorgeaduran/smda-rs.git", branch = "fixes"}
smda = "0.2.6"
thiserror = "1"
walkdir = "2.3.2"
yaml-rust = "0.4.5"
goblin = { version = "0.8.0", features = ["alloc"] }
maplit = "1"
#dnfile = { git = "https://github.com/marirs/dnfile-rs.git", branch = "master" }
dnfile = { git = "https://github.com/jorgeaduran/dnfile-rs.git", branch = "features/optimize_decoding" }
dnfile = { git = "https://github.com/marirs/dnfile-rs.git", branch = "master" }
#dnfile = { git = "https://github.com/jorgeaduran/dnfile-rs.git", branch = "features/optimize_decoding" }
lazy_static = "1.4.0"
parking_lot = "0.12.1"
serde_json = "1.0.113"
Expand Down
12 changes: 6 additions & 6 deletions examples/capa_cli.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ fn main() {
match FileCapabilities::from_file(&filename, &rules_path, true, true, &|_s| {}) {
Err(e) => println!("{:?}", e),
Ok(s) => {
match to_value(&s) {
match to_value(s) {
Err(e) => println!("serde_json_error: {}", e),
Ok(data) => {
let data = data.as_object().unwrap();
Expand Down Expand Up @@ -120,7 +120,7 @@ fn get_properties(props: &Value, features: Option<&Value>) -> Table {
Alignment::CENTER,
)
.with_hspan(2)]));
for (k, v) in &*meta {
for (k, v) in meta {
tbl.add_row(Row::new(vec![
Cell::new(k)
.with_style(Attr::ForegroundColor(color::BRIGHT_BLUE))
Expand Down Expand Up @@ -153,7 +153,7 @@ fn get_mitre(attacks: &Map<String, Value>) -> Table {
Cell::new_align("ATT&CK Technique", Alignment::LEFT),
]));

for (tatic, v) in &*attacks {
for (tatic, v) in attacks {
let techniques = v.as_array().unwrap();
let techniques = techniques
.iter()
Expand Down Expand Up @@ -183,7 +183,7 @@ fn get_mbc(mbc: &Map<String, Value>) -> Table {
Cell::new_align("MBC Objective", Alignment::LEFT),
Cell::new_align("MBC Behavior", Alignment::LEFT),
]));
for (objective, v) in &*mbc {
for (objective, v) in mbc {
let behaviors = v.as_array().unwrap();
let behaviours = behaviors
.iter()
Expand Down Expand Up @@ -213,7 +213,7 @@ fn get_namespace(namespace: &Map<String, Value>) -> Table {
Cell::new_align("Capability", Alignment::LEFT),
Cell::new_align("Namespace", Alignment::LEFT),
]));
for (capability, v) in &*namespace {
for (capability, v) in namespace {
let ns = v.as_str().unwrap().to_string();

tbl.add_row(Row::new(vec![
Expand All @@ -236,7 +236,7 @@ fn get_verbose_info(extra: &Map<String, Value>) -> Table {
Cell::new_align("Features", Alignment::LEFT),
Cell::new_align("Capabilities", Alignment::LEFT),
]));
for (function, v) in &*extra {
for (function, v) in extra {
let caps = v.as_object().unwrap();
let address = caps.get("address").unwrap().as_str().unwrap();
let features = caps.get("features").unwrap().as_u64().unwrap().to_string();
Expand Down
90 changes: 50 additions & 40 deletions src/extractor/dnfile.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ use dnfile::{
};
use std::{
collections::{HashMap, HashSet},
sync::{Arc},
sync::Arc,
};

use parking_lot::RwLock;
Expand Down Expand Up @@ -175,7 +175,7 @@ impl super::Extractor for Extractor {
OpCodeValue::Jmp,
OpCodeValue::Newobj,
]
.contains(&insn.opcode.value)
.contains(&insn.opcode.value)
{
continue;
}
Expand Down Expand Up @@ -210,15 +210,15 @@ impl super::Extractor for Extractor {
) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
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)?,
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
}))
.into_iter()
.fold(Vec::new(), |mut acc, f| {
acc.extend(f);
acc
}))
}

fn get_basic_blocks(
Expand Down Expand Up @@ -260,7 +260,6 @@ impl super::Extractor for Extractor {
ss.extend(self.extract_unmanaged_call_characteristic_features(&f.f, &insn.i)?);
Ok(ss)
}

}

impl Extractor {
Expand Down Expand Up @@ -676,14 +675,14 @@ impl Extractor {
insn: &cil::instruction::Instruction,
) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
let mut res = vec![];
if !vec![
if ![
OpCodeValue::Call,
OpCodeValue::Callvirt,
OpCodeValue::Jmp,
OpCodeValue::Calli,
OpCodeValue::Newobj,
]
.contains(&insn.opcode.value)
.contains(&insn.opcode.value)
{
return Ok(vec![]);
}
Expand Down Expand Up @@ -760,13 +759,13 @@ impl Extractor {
insn: &cil::instruction::Instruction,
) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
let mut res = Vec::new();
if vec![
if [
OpCodeValue::Call,
OpCodeValue::Callvirt,
OpCodeValue::Jmp,
OpCodeValue::Calli,
]
.contains(&insn.opcode.value)
.contains(&insn.opcode.value)
{
let operand_result = resolve_dotnet_token(
&self.pe,
Expand Down Expand Up @@ -795,23 +794,32 @@ impl Extractor {
// Verifica si el nombre del método indica un acceso a una propiedad (get o set).
if operand.name.starts_with("get_") || operand.name.starts_with("set_") {
// Obtiene el namespace y el nombre de la clase a la que pertenece el MemberRef.
let (operand_class_type_namespace, operand_class_type_name) = match operand.class.table() {
"TypeRef" => {
if let Ok(rr) = self.pe.net()?.resolve_coded_index::<TypeRef>(&operand.class) {
(rr.type_namespace.clone(), rr.type_name.clone())
} else {
return Ok(vec![]);
let (operand_class_type_namespace, operand_class_type_name) =
match operand.class.table() {
"TypeRef" => {
if let Ok(rr) = self
.pe
.net()?
.resolve_coded_index::<TypeRef>(&operand.class)
{
(rr.type_namespace.clone(), rr.type_name.clone())
} else {
return Ok(vec![]);
}
}
},
"TypeDef" => {
if let Ok(rr) = self.pe.net()?.resolve_coded_index::<TypeDef>(&operand.class) {
(rr.type_namespace.clone(), rr.type_name.clone())
} else {
return Ok(vec![]);
"TypeDef" => {
if let Ok(rr) = self
.pe
.net()?
.resolve_coded_index::<TypeDef>(&operand.class)
{
(rr.type_namespace.clone(), rr.type_name.clone())
} else {
return Ok(vec![]);
}
}
},
_ => return Ok(vec![]),
};
_ => return Ok(vec![]),
};

// Construye el nombre completo de la propiedad accedida.
let property_name = format!(
Expand Down Expand Up @@ -842,20 +850,22 @@ impl Extractor {
}
}
}
} else if vec![
} else if [
OpCodeValue::Ldfld,
OpCodeValue::Ldflda,
OpCodeValue::Ldsfld,
OpCodeValue::Ldsflda,
OpCodeValue::Stfld,
OpCodeValue::Stsfld,
]
.contains(&insn.opcode.value)
.contains(&insn.opcode.value)
{
if let Ok(fields_lock) = self.get_fields() {
if let Some(fields) = fields_lock.read().as_ref() {
if let Some(field) = fields.get(&(insn.operand.value()? as u64)) {
let access = if vec![OpCodeValue::Stfld, OpCodeValue::Stsfld].contains(&insn.opcode.value) {
let access = if [OpCodeValue::Stfld, OpCodeValue::Stsfld]
.contains(&insn.opcode.value)
{
Some(crate::rules::features::FeatureAccess::Write)
} else {
Some(crate::rules::features::FeatureAccess::Read)
Expand Down Expand Up @@ -932,7 +942,7 @@ impl Extractor {
_f: &cil::function::Function,
insn: &cil::instruction::Instruction,
) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
if !vec![
if ![
OpCodeValue::Call,
OpCodeValue::Callvirt,
OpCodeValue::Jmp,
Expand All @@ -945,7 +955,7 @@ impl Extractor {
OpCodeValue::Stsfld,
OpCodeValue::Newobj,
]
.contains(&insn.opcode.value)
.contains(&insn.opcode.value)
{
return Ok(vec![]);
}
Expand Down Expand Up @@ -983,7 +993,7 @@ impl Extractor {
let fields_lock = self.get_fields()?.read();

if let Some(fields) = &*fields_lock {
if let Some(field) = fields.clone().get(&(insn.operand.value()? as u64)){
if let Some(field) = fields.clone().get(&(insn.operand.value()? as u64)) {
res.push((
crate::rules::features::Feature::Namespace(
crate::rules::features::NamespaceFeature::new(&field.namespace, "")?,
Expand All @@ -1002,7 +1012,7 @@ impl Extractor {
_f: &cil::function::Function,
insn: &cil::instruction::Instruction,
) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
if !vec![
if ![
OpCodeValue::Call,
OpCodeValue::Callvirt,
OpCodeValue::Jmp,
Expand All @@ -1015,7 +1025,7 @@ impl Extractor {
OpCodeValue::Stsfld,
OpCodeValue::Newobj,
]
.contains(&insn.opcode.value)
.contains(&insn.opcode.value)
{
return Ok(vec![]);
}
Expand Down Expand Up @@ -1083,13 +1093,13 @@ impl Extractor {
_f: &cil::function::Function,
insn: &cil::instruction::Instruction,
) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
if !vec![
if ![
OpCodeValue::Call,
OpCodeValue::Callvirt,
OpCodeValue::Jmp,
OpCodeValue::Calli,
]
.contains(&insn.opcode.value)
.contains(&insn.opcode.value)
{
return Ok(vec![]);
}
Expand Down
Loading

0 comments on commit 8fd4803

Please sign in to comment.