phink_lib/cli/ui/
custom.rs

1use crate::{
2    cli::{
3        config::{
4            PFiles,
5            PhinkFiles,
6        },
7        env::PhinkEnv::{
8            AflDebug,
9            AflForkServerTimeout,
10            FromZiggy,
11        },
12        ui::ratatui::CustomUI,
13        ziggy::{
14            ZiggyConfig,
15            AFL_FORKSRV_INIT_TMOUT,
16        },
17    },
18    EmptyResult,
19    ResultOf,
20};
21use anyhow::{
22    bail,
23    Context,
24};
25use std::{
26    fs,
27    process::{
28        Child,
29        Command,
30        Stdio,
31    },
32    sync::mpsc,
33    thread,
34    thread::sleep,
35    time::{
36        Duration,
37        Instant,
38    },
39};
40
41#[derive(Clone, Debug)]
42pub struct CustomManager {
43    args: Option<Vec<String>>,
44    env: Vec<(String, String)>,
45    ziggy_config: ZiggyConfig,
46}
47
48impl CustomManager {
49    pub fn new(
50        args: Option<Vec<String>>,
51        env: Vec<(String, String)>,
52        ziggy_config: ZiggyConfig,
53    ) -> Self {
54        // If it exists, we remove the `PFiles::LastSeed`
55        let _ = fs::remove_file(
56            PhinkFiles::new(ziggy_config.config().fuzz_output.clone().unwrap())
57                .path(PFiles::LastSeed),
58        );
59
60        Self {
61            args,
62            env,
63            ziggy_config,
64        }
65    }
66
67    pub fn cmd_fuzz(self) -> ResultOf<Child> {
68        let mut binding = Command::new("cargo");
69        let command_builder = binding
70            .arg("ziggy")
71            .arg("fuzz")
72            .env(FromZiggy.to_string(), "1")
73            .env(AflForkServerTimeout.to_string(), AFL_FORKSRV_INIT_TMOUT)
74            .env(AflDebug.to_string(), self.ziggy_config.afl_debug())
75            .stdout(Stdio::null())
76            .stderr(Stdio::null());
77
78        self.ziggy_config
79            .with_allowlist(command_builder)
80            .context("Couldn't use the allowlist")?;
81
82        if let Some(args) = self.args {
83            command_builder.args(args.iter());
84        }
85        command_builder.envs(self.env);
86
87        let child = command_builder
88            .spawn()
89            .context("Spawning Ziggy was unsuccessful")?;
90
91        Ok(child)
92    }
93
94    pub fn start(self) -> EmptyResult {
95        let (tx, rx) = mpsc::channel();
96        let cloned_config = self.ziggy_config.clone();
97
98        thread::spawn(move || {
99            let fuzzer_pid = self.cmd_fuzz();
100            tx.send(fuzzer_pid).unwrap();
101        });
102
103        let child: Child = rx.recv()??;
104
105        let mut ratatui = CustomUI::new(&cloned_config);
106        let start_time = Instant::now();
107        let mut print_err = true;
108        loop {
109            if start_time.elapsed() > Duration::new(120, 0) {
110                bail!("Couldn't instantiate the custom UI within 120 seconds...");
111            }
112            if ratatui.is_err() {
113                if print_err {
114                    println!("⏰ Waiting for AFL++ to finish the dry run");
115                    print_err = false;
116                }
117                ratatui = CustomUI::new(&cloned_config);
118                sleep(Duration::from_millis(100));
119            } else {
120                break;
121            }
122        }
123
124        ratatui
125            .context("Couldn't create ratatui with 'new'")?
126            .initialize_tui(child)
127            .context("Couldn't initialize ratatui")?;
128        Ok(())
129    }
130}
phink_lib/cli/ui/custom.rs

phink_lib/cli/ui/
custom.rs
