use std::path::{Path, PathBuf};
use std::str::FromStr;
pub use crate::values::Analysis;
use crate::values::{AnalysisParser, BirthmarkTypeParser};
use crate::vocabulary::{algorithm_parser, ir_parser, strategy_parser};
use clap::ValueEnum;
use oinkie::birthmarks::{AnalysisType, BirthmarkType};
use oinkie::compare::{Aggregator, Algorithm, Comparator, PairingStrategy};
use oinkie::extract::Extractor;
use oinkie::lift::Ir;
use oinkie::{Oinkie, Result};
use std::num::NonZeroUsize;
#[derive(Debug, clap::Parser)]
#[command(version, about, max_term_width = 100)]
pub struct OinkieOpts {
#[clap(subcommand)]
pub command: OinkieCommand,
#[clap(short, long, value_enum, default_value_t = LogLevel::Warn, value_name = "LEVEL", ignore_case = true, help = "Log level for the application")]
pub level: LogLevel,
#[clap(
long,
help = "Draw no progress bars, as when stderr is kept in a log or a CI output"
)]
pub no_progress: bool,
}
fn filter_level(level: &LogLevel) -> log::LevelFilter {
match level {
LogLevel::Debug => log::LevelFilter::Debug,
LogLevel::Info => log::LevelFilter::Info,
LogLevel::Warn => log::LevelFilter::Warn,
LogLevel::Error => log::LevelFilter::Error,
LogLevel::Trace => log::LevelFilter::Trace,
LogLevel::Off => log::LevelFilter::Off,
}
}
impl OinkieOpts {
pub fn init(&self) -> Result<()> {
let filter = filter_level(&self.level);
env_logger::Builder::new().filter_level(filter).init();
Ok(())
}
}
#[derive(Debug, clap::Parser, ValueEnum, Clone)]
pub enum LogLevel {
Error,
Warn,
Info,
Debug,
Trace,
Off,
}
#[derive(Debug, clap::Parser)]
pub enum OinkieCommand {
#[command(name = "info", about = "Display information about the application")]
Info,
#[command(
name = "lift",
about = "Lift binary files to JSON files of an intermediate representation, using a specified lifter"
)]
Lift(LiftOpts),
#[command(
name = "extract",
about = "Extract birthmarks from a lifted binary file (JSON format)"
)]
Extract(ExtractOpts),
#[command(
name = "compare",
about = "Compare birthmarks and output the similarity score"
)]
Compare(CompareOpts),
#[command(
name = "review",
about = "Re-read a finished comparison: recompute the similarity of each pair from the stored function similarities"
)]
Review(ReviewOpts),
#[command(
name = "reaggregate",
hide = true,
disable_help_flag = true,
trailing_var_arg = true
)]
Reaggregate {
#[clap(allow_hyphen_values = true, num_args = 0..)]
args: Vec<String>,
},
#[command(
name = "stats",
about = "Summarise a set of birthmarks: how many of each type, how many functions each holds, and how long each function's birthmark is"
)]
Stats(StatsOpts),
#[command(
name = "run",
about = "Extract birthmarks and compare them in one command"
)]
Run(RunOpts),
#[cfg(feature = "mcp")]
#[command(
name = "mcp",
about = "Serve oinkie over the Model Context Protocol, on stdin and stdout"
)]
Mcp(McpOpts),
}
#[derive(Debug, Clone, clap::Args)]
pub struct Threads {
#[clap(
long,
value_name = "N",
help = "Number of threads to compute on [default: one per core]"
)]
threads: Option<NonZeroUsize>,
}
impl Threads {
pub fn oinkie(&self) -> Result<Oinkie> {
match self.threads {
Some(n) => Oinkie::builder().threads(n).build(),
None => Oinkie::new(),
}
}
}
#[cfg(feature = "mcp")]
#[derive(Debug, clap::Parser)]
pub struct McpOpts {
#[clap(
short = 'r',
long = "root",
value_name = "DIRECTORY",
help = "Directory the tools may read and write under. May be given more than once",
long_help = "Directory the tools may read and write under. May be given more than once. \
Every path a tool is handed, input and output alike, has to resolve inside one of these; anything \
else is refused. Defaults to the working directory. The paths the tools receive are written by a \
language model rather than by you, which is the whole reason this exists."
)]
roots: Vec<PathBuf>,
#[clap(flatten)]
pub(crate) threads: Threads,
}
#[cfg(feature = "mcp")]
impl McpOpts {
pub fn roots(&self) -> &[PathBuf] {
&self.roots
}
}
#[derive(Debug, clap::Parser)]
pub struct LiftOpts {
#[clap(
short,
long,
default_value = "pcodes",
value_name = "DIRECTORY",
help = "Specify the directory for putting the resultant JSON files of the lifted programs (default: './pcodes' directory)"
)]
dest: PathBuf,
#[clap(
short = 'r',
long,
value_parser = ir_parser(),
default_value_t = Ir::default(),
help = "Intermediate representation to produce. This also picks the tool, since a representation is only produced by one of them; several representations can come from the same tool."
)]
ir: Ir,
#[clap(
short = 'H',
long,
value_name = "HOME",
help = "Path to the installation directory of the tool behind --ir. If not specified, that tool's own environment variable (GHIDRA_HOME for Ghidra) is read, then the usual install locations are searched. The error names which variable to set."
)]
home: Option<PathBuf>,
#[clap(
short = 'i',
long = "intermediate",
value_name = "DIRECTORY",
help = "Directory for the lifter to work in, kept rather than discarded. Every lifter runs in one, since that is where its script writes; Ghidra also keeps its project there. If not specified, a temporary directory is used and deleted."
)]
intermediate_dir: Option<PathBuf>,
#[clap(
long,
value_name = "SCRIPT",
help = "Path to a custom lifting script, replacing the built-in one. The language is that of the tool behind --ir: Java for Ghidra. It must write {input file name}.json into its working directory."
)]
script: Option<PathBuf>,
#[clap(
short = 'j',
long,
default_value = "1",
value_name = "N",
value_parser = parse_jobs,
help = "Lift up to N files at a time (default: 1, one after another). Lifting runs a whole decompiler process per file, and several of them against a Ghidra installation whose language cache has not been built yet can corrupt it, so parallelism is opt-in."
)]
jobs: std::num::NonZeroUsize,
#[clap(
short = 'S',
long,
default_value_t = false,
help = "Skip if the resultant JSON file already exists"
)]
skip: bool,
#[clap(
index = 1,
value_name = "FILES",
help = "Path to the binary or intermediate files to lift"
)]
files: Vec<PathBuf>,
}
fn parse_jobs(s: &str) -> std::result::Result<std::num::NonZeroUsize, String> {
match s.parse::<usize>() {
Ok(n) => std::num::NonZeroUsize::new(n).ok_or_else(|| "must be at least 1".to_string()),
Err(e) => Err(e.to_string()),
}
}
impl LiftOpts {
pub fn dest(&self) -> &Path {
&self.dest
}
pub fn ir(&self) -> Ir {
self.ir
}
pub fn home(&self) -> Option<&Path> {
self.home.as_deref()
}
pub fn intermediate_dir(&self) -> Option<&Path> {
self.intermediate_dir.as_deref()
}
pub fn script(&self) -> Option<&Path> {
self.script.as_deref()
}
pub fn is_skip(&self) -> bool {
self.skip
}
pub fn jobs(&self) -> usize {
self.jobs.get()
}
pub fn iter(&self) -> impl Iterator<Item = &PathBuf> {
self.files.iter()
}
pub fn len(&self) -> usize {
self.files.len()
}
}
#[derive(Debug, clap::Parser)]
pub struct ExtractOpts {
#[clap(
short,
long,
default_value = "birthmarks",
value_name = "DIRECTORY",
help = "Specify the directory for putting the resultant JSON files for the extracted birthmarks (default: './birthmarks' directory)"
)]
dest: PathBuf,
#[clap(short, long, value_name = "BIRTHMARK_TYPE", value_parser = BirthmarkTypeParser, default_value = "op-seq", hide_possible_values = true, help = "Type of birthmark to extract, such as 'op-seq' or 'fc-freq'; 'oinkie info' lists them", long_help = "Type of birthmark to extract. \
fc (Function Calls) and op (Opcode) with set, seq, and freq variants are supported. \
For example, 'op-seq' extracts the sequence of operations as a birthmark, \
while 'fc-freq' extracts the frequency of function calls. \
k-grams are written with the k in the name: 'op-3gram-set'. Any k parses, not \
only the ones 'oinkie info' lists. \
The full birthmark types can be found by running 'oinkie info'.")]
birthmark_type: BirthmarkType,
#[clap(
short = 'S',
long,
default_value_t = false,
help = "Skip the resultant birthmark file is already exists"
)]
skip: bool,
#[clap(
index = 1,
value_name = "JSON_FILES",
help = "Path to the JSON files to extract birthmarks from"
)]
files: Vec<PathBuf>,
#[clap(flatten)]
pub(crate) threads: Threads,
}
impl ExtractOpts {
pub fn dest(&self) -> &Path {
&self.dest
}
pub fn extractor(&self) -> Extractor {
Extractor::new(self.birthmark_type.clone())
}
pub fn len(&self) -> usize {
self.files.len()
}
pub fn is_empty(&self) -> bool {
self.files.is_empty()
}
pub fn iter(&self) -> impl Iterator<Item = &PathBuf> {
self.files.iter()
}
pub fn is_skip(&self) -> bool {
self.skip
}
}
#[derive(Debug, clap::Parser)]
pub struct ReviewOpts {
#[clap(
short = 'A',
long,
default_value = "hungarian",
value_name = "METHOD",
ignore_case = true,
help = "Aggregator combining a pair's function similarities into its similarity: hungarian, topn:N, containment, matched:T or weighted",
long_help = "Specify the aggregator for combining the function similarities of a pair into its similarity.
- hungarian: use the Hungarian algorithm to find the one-to-one matching between the functions of the two birthmarks that maximises the total similarity.
- topn:N: take each function's best match in the other birthmark, and average the N best of those from each side. Matches need not be one-to-one, which can reduce noise from less relevant matches. topn:all, or topn, takes them all.
- containment: as hungarian, but average over the smaller birthmark's functions rather than the larger's: how much of the smaller is found in the larger. A program copied whole into a larger one scores 1.0. It says the smaller is contained, not that the two are alike.
- matched:T: as hungarian, but the proportion of the larger birthmark's functions matched with a similarity of at least T, where 0 < T <= 1.
- weighted: as hungarian, but each function counts by its number of elements, so that short functions, which agree by chance, count for little. review reads the sizes from the birthmarks the comparison names."
)]
aggregator: Aggregator,
#[clap(
short,
long,
value_name = "RESULT.CSV",
help = "Specify the result CSV file of the comparing results to review. \
The file lists the similarity of each pair.",
default_value = "review.csv"
)]
dest_file: PathBuf,
#[clap(
index = 1,
value_name = "SCORE_DIRECTORY",
help = "Path to the score directory: the pair CSVs that compare or run wrote, holding the function similarities"
)]
score_directory: PathBuf,
#[clap(
long,
value_name = "N|Rx",
help = "Drop the functions with fewer elements than this before aggregating: N elements, or R times the mean (Rx)",
long_help = "Drop the functions with fewer elements than this before aggregating. \
N is a count of elements; Rx is R times the mean count over every function of every birthmark \
in the directory (e.g. 0.3x). Needs the birthmarks the comparisons name. The threshold is \
recorded on the last line of the summary."
)]
min_elements: Option<MinElements>,
#[clap(flatten)]
pub(crate) threads: Threads,
}
impl ReviewOpts {
pub fn aggregator(&self) -> &Aggregator {
&self.aggregator
}
pub fn score_directory(&self) -> &Path {
&self.score_directory
}
pub fn dest_file(&self) -> &PathBuf {
&self.dest_file
}
pub fn min_elements(&self) -> Option<&crate::cli::MinElements> {
self.min_elements.as_ref()
}
}
#[derive(Debug, Clone, PartialEq)]
pub enum MinElements {
Count(usize),
Ratio(f64),
}
impl FromStr for MinElements {
type Err = String;
fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
let s = s.trim();
if let Some(ratio) = s.strip_suffix('x') {
return match ratio.parse::<f64>() {
Ok(r) if r.is_finite() && r >= 0.0 => Ok(MinElements::Ratio(r)),
_ => Err(format!(
"{s}: a ratio of the mean is a number of zero or more followed by x, such as 0.3x"
)),
};
}
match s.parse::<usize>() {
Ok(n) => Ok(MinElements::Count(n)),
Err(_) if s.parse::<f64>().is_ok() => Err(format!(
"{s}: a count of elements is a whole number; for {s} times the mean, write {s}x"
)),
Err(_) => Err(format!(
"{s}: expected a count of elements, such as 5, or a ratio of the mean, such as 0.3x"
)),
}
}
}
impl std::fmt::Display for MinElements {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
MinElements::Count(n) => write!(f, "{n}"),
MinElements::Ratio(r) => write!(f, "{r}x"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
pub enum StatsFormat {
Json,
Csv,
Markdown,
}
#[derive(Debug, clap::Parser)]
pub struct StatsOpts {
#[clap(
short,
long,
value_enum,
default_value_t = StatsFormat::Markdown,
value_name = "FORMAT",
ignore_case = true,
help = "Output format"
)]
format: StatsFormat,
#[clap(
short,
long,
value_name = "FILE",
help = "Write the statistics to FILE rather than to standard output"
)]
output: Option<PathBuf>,
#[clap(
short,
long,
default_value_t = false,
help = "Descend into the subdirectories of the given directories"
)]
recursive: bool,
#[clap(
long,
default_value_t = false,
help = "Report each birthmark file as well as each group. With -f csv this replaces the summary table."
)]
per_file: bool,
#[clap(
short,
long,
value_name = "N",
help = "Report the N most frequent elements of each group. With -f csv this replaces the summary table."
)]
top: Option<usize>,
#[clap(
index = 1,
required = true,
value_name = "PATHS",
help = "Birthmark files, or directories holding them. A directory contributes the *.json directly inside it; \
a file that does not read as a birthmark is skipped with a warning and counted."
)]
inputs: Vec<PathBuf>,
#[clap(flatten)]
pub(crate) threads: Threads,
}
impl StatsOpts {
pub fn format(&self) -> StatsFormat {
self.format
}
pub fn dest(&self) -> Option<&Path> {
self.output.as_deref()
}
pub fn is_recursive(&self) -> bool {
self.recursive
}
pub fn is_per_file(&self) -> bool {
self.per_file
}
pub fn top(&self) -> Option<usize> {
self.top
}
pub fn inputs(&self) -> &[PathBuf] {
&self.inputs
}
}
#[derive(Debug, clap::Parser)]
pub struct CompareOpts {
#[clap(short, long, value_parser = algorithm_parser(), default_value = "jaccard", value_name = "ALGORITHM", ignore_case = true, help = "Specify the similarity calculation algorithm.")]
algorithm: Algorithm,
#[clap(
short = 'A',
long,
default_value = "hungarian",
value_name = "METHOD",
ignore_case = true,
help = "Aggregator combining a pair's function similarities into its similarity: hungarian, topn:N, containment, matched:T or weighted",
long_help = "Specify the aggregator for combining the function similarities of a pair into its similarity.
- hungarian: use the Hungarian algorithm to find the one-to-one matching between the functions of the two birthmarks that maximises the total similarity.
- topn:N: take each function's best match in the other birthmark, and average the N best of those from each side. Matches need not be one-to-one, which can reduce noise from less relevant matches. topn:all, or topn, takes them all.
- containment: as hungarian, but average over the smaller birthmark's functions rather than the larger's: how much of the smaller is found in the larger. A program copied whole into a larger one scores 1.0. It says the smaller is contained, not that the two are alike.
- matched:T: as hungarian, but the proportion of the larger birthmark's functions matched with a similarity of at least T, where 0 < T <= 1.
- weighted: as hungarian, but each function counts by its number of elements, so that short functions, which agree by chance, count for little. review reads the sizes from the birthmarks the comparison names."
)]
aggregator: Aggregator,
#[clap(short, long, value_parser = strategy_parser(), default_value = "all-and-self", value_name = "STRATEGY", ignore_case = true, help = "Pairing strategy for comparing files")]
strategy: PairingStrategy,
#[clap(
short,
long,
value_name = "DIRECTORY",
help = "Destination directory for the results",
default_value = "similarities"
)]
dest: PathBuf,
#[clap(
short = 'S',
long,
default_value_t = false,
help = "Skip if the similarity file already exists for the pair of birthmarks"
)]
skip: bool,
#[clap(
index = 1,
value_name = "JSON_FILES",
help = "Path to the birthmark JSON files to compare"
)]
files: Vec<PathBuf>,
#[clap(flatten)]
pub(crate) threads: Threads,
}
impl CompareOpts {
pub fn dest(&self) -> &Path {
&self.dest
}
pub fn comparator(&self, oinkie: &Oinkie) -> Comparator {
oinkie.comparator(&self.algorithm)
}
pub fn iter(&self) -> Box<dyn Iterator<Item = (&PathBuf, &PathBuf)> + Send + '_> {
self.strategy.pairs(&self.files)
}
pub fn aggregator(&self) -> &Aggregator {
log::info!(
"Using {:?} as the aggregator for combining function similarities",
self.aggregator
);
&self.aggregator
}
pub fn compare_count(&self) -> usize {
self.strategy.compare_count(&self.files)
}
pub fn is_skip(&self) -> bool {
self.skip
}
}
#[derive(Debug, clap::Parser)]
pub struct RunOpts {
#[clap(short, long, value_name = "ANALYSIS", value_parser = AnalysisParser, default_value = "op-set-jaccard", hide_possible_values = true, help = "Analysis to run, as '{birthmark}-{algorithm}' -- for example 'op-set-jaccard' or 'op-3gram-freq-cosine'", long_help = "Analysis to run, as '{birthmark}-{algorithm}' -- for example 'op-set-jaccard' or 'op-3gram-freq-cosine'. \
Run 'oinkie info' for the birthmarks and the algorithms they pair with. Any k \
parses in a k-gram name, not only the ones listed.")]
pub(crate) analysis: Analysis,
#[clap(short, long, value_parser = strategy_parser(), default_value = "all-and-self", ignore_case = true, help = "Pairing strategy for file comparisons")]
pub strategy: PairingStrategy,
#[clap(
short,
long,
default_value = "similarities",
help = "Destination path for the output CSV file (default: 'similarities' directory"
)]
pub(crate) dest: PathBuf,
#[clap(
short = 'A',
long,
default_value = "hungarian",
value_name = "METHOD",
ignore_case = true,
help = "Aggregator combining a pair's function similarities into its similarity: hungarian, topn:N, containment, matched:T or weighted",
long_help = "Specify the aggregator for combining the function similarities of a pair into its similarity.
- hungarian: use the Hungarian algorithm to find the one-to-one matching between the functions of the two birthmarks that maximises the total similarity.
- topn:N: take each function's best match in the other birthmark, and average the N best of those from each side. Matches need not be one-to-one, which can reduce noise from less relevant matches. topn:all, or topn, takes them all.
- containment: as hungarian, but average over the smaller birthmark's functions rather than the larger's: how much of the smaller is found in the larger. A program copied whole into a larger one scores 1.0. It says the smaller is contained, not that the two are alike.
- matched:T: as hungarian, but the proportion of the larger birthmark's functions matched with a similarity of at least T, where 0 < T <= 1.
- weighted: as hungarian, but each function counts by its number of elements, so that short functions, which agree by chance, count for little. review reads the sizes from the birthmarks the comparison names."
)]
aggregator: Aggregator,
#[clap(
short = 'S',
long,
default_value_t = false,
help = "Skip if the similarity file already exists for the pair of birthmarks"
)]
pub(crate) skip: bool,
#[clap(index = 1, help = "Path to the JSON files")]
pub(crate) files: Vec<PathBuf>,
#[clap(flatten)]
pub(crate) threads: Threads,
}
impl RunOpts {
pub fn dest(&self) -> &Path {
&self.dest
}
pub fn analysis_type(&self) -> Result<AnalysisType> {
self.analysis.analysis_type()
}
pub fn iter(&self) -> Box<dyn Iterator<Item = (&PathBuf, &PathBuf)> + Send + '_> {
self.strategy.pairs(&self.files)
}
pub fn compare_count(&self) -> usize {
self.strategy.compare_count(&self.files)
}
pub fn is_skip(&self) -> bool {
self.skip
}
pub fn aggregator(&self) -> &Aggregator {
log::info!(
"Using {:?} as the aggregator for combining function similarities",
self.aggregator
);
&self.aggregator
}
}
#[cfg(test)]
mod tests {
use super::*;
use clap::Parser;
#[test]
fn test_long_help_opens_with_the_commands_own_description() {
use clap::CommandFactory;
let mut root = OinkieOpts::command();
root.build();
for sub in root.get_subcommands_mut().filter(|s| !s.is_hide_set()) {
let name = sub.get_name().to_string();
let Some(about) = sub.get_about().map(ToString::to_string) else {
continue;
};
let help = sub.render_long_help().to_string();
let opening = help
.lines()
.take_while(|l| !l.trim().is_empty())
.collect::<Vec<_>>()
.join(" ");
assert_eq!(
opening, about,
"`oinkie {name} --help` opens with something else:\n{help}"
);
}
}
#[test]
fn test_every_log_level_maps_to_a_filter() {
let cases = [
(LogLevel::Trace, log::LevelFilter::Trace),
(LogLevel::Debug, log::LevelFilter::Debug),
(LogLevel::Info, log::LevelFilter::Info),
(LogLevel::Warn, log::LevelFilter::Warn),
(LogLevel::Error, log::LevelFilter::Error),
(LogLevel::Off, log::LevelFilter::Off),
];
for (level, expected) in cases {
assert_eq!(filter_level(&level), expected, "{level:?}");
}
}
#[test]
fn test_jobs_refuses_what_is_not_a_count() {
assert_eq!(parse_jobs("4").unwrap().get(), 4);
assert_eq!(parse_jobs("0").unwrap_err(), "must be at least 1");
let e = parse_jobs("many").unwrap_err();
assert!(e.contains("invalid digit"), "{e}");
}
#[test]
fn test_run_takes_skip_and_reports_it() {
for (args, expected) in [
(vec!["oinkie", "run", "a.json"], false),
(vec!["oinkie", "run", "-S", "a.json"], true),
] {
let OinkieCommand::Run(opts) = OinkieOpts::try_parse_from(args).unwrap().command else {
panic!("Expected Run command");
};
assert_eq!(opts.is_skip(), expected);
}
}
}