sigstrike 0.1.4

Cobalt Strike beacon crawler and parser.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158

use crate::io::{list_files, process_files};
use clap::{Parser, Subcommand};
use env_logger::Env;
use log::{debug, error, info};
use std::path::PathBuf;
use std::time::Instant;

#[cfg(any(target_os = "linux", target_os = "macos"))]
use libc::{RLIMIT_NOFILE, rlimit, setrlimit};

#[cfg(any(target_os = "linux", target_os = "macos"))]
fn set_nofile_limit(soft: u64, hard: u64) -> Result<(), String> {
    let lim = rlimit {
        rlim_cur: soft,
        rlim_max: hard,
    };
    let res = unsafe { setrlimit(RLIMIT_NOFILE, &lim) };
    if res == 0 {
        Ok(())
    } else {
        Err(std::io::Error::last_os_error().to_string())
    }
}

#[derive(Parser)]
#[command(name = "sigstrike")]
#[command(about = "An example app with subcommands", long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Commands,

    #[arg(
        long,
        default_value = "error",
        help = "Logging level (trace, debug, info, warn, error)",
        global = true
    )]
    logging_level: String,
}

#[derive(Subcommand)]
enum Commands {
    Process {
        #[arg(
            long,
            help = "Input path to a file or directory containing Cobalt Strike beacons."
        )]
        input_path: PathBuf,

        #[arg(
            long,
            help = "Output path for results. If not specified, results will be printed to stdout."
        )]
        output_path: Option<PathBuf>,
    },
    Crawl {
        #[arg(long, help = "Path to input file containing URLs")]
        input_path: PathBuf,

        #[arg(long, help = "Path to output file (JSONL format)")]
        output_path: PathBuf,

        #[arg(long, default_value_t = 100)]
        max_concurrent: usize,

        #[arg(
            long,
            default_value_t = 2,
            help = "Maximum number of retries for each URL"
        )]
        max_retries: usize,

        #[arg(
            long,
            default_value_t = 10,
            help = "Timeout in seconds for each request"
        )]
        timeout: u64,
    },
}

pub fn parse_beacons(input_path: PathBuf, output_path: Option<PathBuf>) {
    let input_str_path = input_path.display().to_string();

    if input_path.is_file() {
        let files: Vec<String> = Vec::from([input_str_path.clone()]);
        if let Err(e) = process_files(files, output_path) {
            error!("Error processing file {}: {}", &input_str_path, e);
        }
    } else if input_path.is_dir() {
        let files = list_files(&input_str_path);
        info!(
            "Found {} files in directory: {}",
            files.len(),
            &input_str_path
        );
        if files.is_empty() {
            error!("No files found in directory: {input_str_path}");
        }
        if let Err(e) = process_files(files, output_path) {
            error!(
                "Error processing files in directory {}: {}",
                &input_str_path, e
            );
        }
    } else {
        error!(
            "Input path is neither a file nor a directory: {}",
            &input_str_path
        );
    };
}

pub async fn run_cli(start_arg: usize) {
    let start = Instant::now();
    let cli = Cli::parse_from(std::env::args().skip(start_arg));

    let env = Env::default().default_filter_or(cli.logging_level);
    env_logger::Builder::from_env(env).init();

    #[cfg(any(target_os = "linux", target_os = "macos"))]
    match set_nofile_limit(65535, 65535) {
        Ok(()) => debug!("File limit set"),
        Err(e) => error!("Failed to set limit: {e}"),
    }

    match cli.command {
        Commands::Process {
            input_path,
            output_path,
        } => {
            parse_beacons(input_path, output_path);
        }
        Commands::Crawl {
            input_path,
            output_path,
            max_concurrent,
            max_retries,
            timeout,
        } => {
            let crawl_result = crate::crawler::crawl(
                &input_path,
                &output_path,
                max_concurrent,
                max_retries,
                timeout,
            )
            .await;
            if let Err(e) = crawl_result {
                error!("Error running crawler: {e}");
            }
        }
    }

    let end = Instant::now();
    let duration = end.duration_since(start);
    info!("Total execution time: {duration:?}");
}