diff --git a/README.md b/README.md index 9f18dbd..646f77e 100644 --- a/README.md +++ b/README.md @@ -30,6 +30,8 @@ Run `pwninit` in a directory with the relevant files and it will detect which on If you don't like the default template, you can use your own. Just specify `--template-path `. Check [template.py](src/template.py) for the template format. The names of the `exe`, `libc`, and `ld` bindings can be customized with `--template-bin-name`, `--template-libc-name`, and `--template-ld-name`. +The name of the solve script can be configured with `--template-name `. For example: `--template-name exploit.py` will generate a solve script with the name `exploit.py`. + ##### Persisting custom `solve.py` You can make `pwninit` load your custom template automatically by adding an alias to your `~/.bashrc`. diff --git a/src/opts.rs b/src/opts.rs index a6be341..65caf05 100644 --- a/src/opts.rs +++ b/src/opts.rs @@ -66,6 +66,11 @@ pub struct Opts { #[structopt(default_value = "libc")] pub template_libc_name: String, + /// Name of pwntools solve script + #[structopt(long)] + #[structopt(default_value = "solve.py")] + pub template_name: String, + /// Name of linker variable for pwntools solve script #[structopt(long)] #[structopt(default_value = "ld")] diff --git a/src/solvepy.rs b/src/solvepy.rs index 576ae3f..a134d6d 100644 --- a/src/solvepy.rs +++ b/src/solvepy.rs @@ -87,13 +87,16 @@ fn make_stub(opts: &Opts) -> Result { .context(FmtSnafu) } -/// Write script produced with `make_stub()` to `solve.py` in the -/// specified directory, unless a `solve.py` already exists +/// Write script produced with `make_stub()` to `opts.template_name` in the +/// specified directory, unless a `opts.template_name` already exists pub fn write_stub(opts: &Opts) -> Result<()> { let stub = make_stub(opts)?; - let path = Path::new("solve.py"); + let path = Path::new(&opts.template_name); if !path.exists() { - println!("{}", "writing solve.py stub".cyan().bold()); + println!( + "{}", + format!("writing {} stub", opts.template_name).cyan().bold() + ); fs::write(path, stub).context(WriteSnafu)?; set_exec(path).context(SetExecSnafu)?; }