Skip to content

Commit

Permalink
Merge pull request #8 from jorgeaduran/master
Browse files Browse the repository at this point in the history 
Enhanced Feature Extraction and Output Customization in Capa CLI
  • Loading branch information
@marirs
marirs authored Feb 15, 2024
2 parents 87400ed + 2f40244 commit a2b0399
Show file tree
Hide file tree
Showing 4 changed files with 244 additions and 393 deletions.
34 changes: 24 additions & 10 deletions examples/capa_cli.rs
View file
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use std::fs;
use capa::FileCapabilities;
use clap::Parser;
use prettytable::{color, format::Alignment, Attr, Cell, Row, Table};
Expand All @@ -6,10 +7,10 @@ use std::time::Instant;

#[derive(Parser)]
#[clap(
author,
version,
about,
long_about = "Find Capabilities of a given file!"
author,
version,
about,
long_about = "Find Capabilities of a given file!"
)]
struct CliOpts {
/// File to analyse
Expand All @@ -21,19 +22,27 @@ struct CliOpts {
/// verbose output
#[clap(long)]
verbose: bool,
/// file path to save the result in json format
#[clap(short = 'o', long, value_name = "JSON_PATH")]
output: Option<String>,
/// map_features
#[clap(short = 'm', long, default_value = "false")]
map_features: bool,
}

fn main() {
let cli = CliOpts::parse();
let filename = cli.name;
let rules_path = cli.rules_path;
let verbose = cli.verbose;
let map_features = cli.map_features;
let json_path = cli.output;

let start = Instant::now();
match FileCapabilities::from_file(&filename, &rules_path, true, true, &|_s| {}) {
match FileCapabilities::from_file(&filename, &rules_path, true, true, &|_s| {}, map_features) {
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 @@ -105,6 +114,11 @@ fn main() {
println!();
}
}
if let Some(json_path) = json_path {
let json = s.serialize_file_capabilities().unwrap();
fs::write(json_path.clone(), json).expect("Unable to write file");
println!("Analysis result saved in JSON format at: {}", json_path);
}
}
}
println!("Time taken (seconds): {:?}", start.elapsed());
Expand All @@ -119,7 +133,7 @@ fn get_properties(props: &Value, features: Option<&Value>) -> Table {
"File Properties",
Alignment::CENTER,
)
.with_hspan(2)]));
.with_hspan(2)]));
for (k, v) in meta {
tbl.add_row(Row::new(vec![
Cell::new(k)
Expand Down Expand Up @@ -147,7 +161,7 @@ fn get_mitre(attacks: &Map<String, Value>) -> Table {
"MITRE ATT&CK",
Alignment::CENTER,
)
.with_hspan(2)]));
.with_hspan(2)]));
tbl.set_titles(Row::new(vec![
Cell::new_align("ATT&CK Tactic", Alignment::LEFT),
Cell::new_align("ATT&CK Technique", Alignment::LEFT),
Expand Down Expand Up @@ -178,7 +192,7 @@ fn get_mbc(mbc: &Map<String, Value>) -> Table {
"Malware Behavior Catalog",
Alignment::CENTER,
)
.with_hspan(2)]));
.with_hspan(2)]));
tbl.set_titles(Row::new(vec![
Cell::new_align("MBC Objective", Alignment::LEFT),
Cell::new_align("MBC Behavior", Alignment::LEFT),
Expand Down Expand Up @@ -208,7 +222,7 @@ fn get_namespace(namespace: &Map<String, Value>) -> Table {
"File Capability/Namespace",
Alignment::CENTER,
)
.with_hspan(2)]));
.with_hspan(2)]));
tbl.set_titles(Row::new(vec![
Cell::new_align("Capability", Alignment::LEFT),
Cell::new_align("Namespace", Alignment::LEFT),
Expand Down
30 changes: 15 additions & 15 deletions src/extractor/dnfile.rs
View file
Original file line number Diff line number Diff line change
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 @@ -682,7 +682,7 @@ impl Extractor {
OpCodeValue::Calli,
OpCodeValue::Newobj,
]
.contains(&insn.opcode.value)
.contains(&insn.opcode.value)
{
return Ok(vec![]);
}
Expand Down Expand Up @@ -765,7 +765,7 @@ impl Extractor {
OpCodeValue::Jmp,
OpCodeValue::Calli,
]
.contains(&insn.opcode.value)
.contains(&insn.opcode.value)
{
let operand_result = resolve_dotnet_token(
&self.pe,
Expand Down Expand Up @@ -858,7 +858,7 @@ impl Extractor {
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() {
Expand Down Expand Up @@ -955,7 +955,7 @@ impl Extractor {
OpCodeValue::Stsfld,
OpCodeValue::Newobj,
]
.contains(&insn.opcode.value)
.contains(&insn.opcode.value)
{
return Ok(vec![]);
}
Expand Down Expand Up @@ -1025,7 +1025,7 @@ impl Extractor {
OpCodeValue::Stsfld,
OpCodeValue::Newobj,
]
.contains(&insn.opcode.value)
.contains(&insn.opcode.value)
{
return Ok(vec![]);
}
Expand Down Expand Up @@ -1099,7 +1099,7 @@ impl Extractor {
OpCodeValue::Jmp,
OpCodeValue::Calli,
]
.contains(&insn.opcode.value)
.contains(&insn.opcode.value)
{
return Ok(vec![]);
}
Expand Down
6 changes: 3 additions & 3 deletions src/extractor/smda.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ impl super::Extractor for Extractor {
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_embedded_pe()?);
res.extend(self.extract_file_embedded_pe()?);
res.extend(self.extract_file_strings()?);
// res.extend(self.extract_file_function_names(pbytes)?);
res.extend(self.extract_file_format()?);
Expand Down Expand Up @@ -334,10 +334,10 @@ impl Extractor {
Ok(res)
}

fn _extract_file_embedded_pe(&self) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
fn extract_file_embedded_pe(&self) -> Result<Vec<(crate::rules::features::Feature, u64)>> {
let mut res = vec![];
for (mz_offset, _pe_offset, _key) in
Extractor::find_embedded_pe_headers(&self.report.buffer)
Extractor::find_embedded_pe_headers(&self.report.buffer)
{
res.push((
crate::rules::features::Feature::Characteristic(
Expand Down
Loading

0 comments on commit a2b0399

Please sign in to comment.