Skip to content

Commit

Permalink
Split up code into modules; use reintroduce static dispatch for `Patt…
Browse files Browse the repository at this point in the history 
…erns`
  • Loading branch information
@Limeth
Limeth committed Feb 25, 2018
1 parent 62e75a3 commit 050f601
Show file tree
Hide file tree
Showing 5 changed files with 257 additions and 227 deletions.
25 changes: 25 additions & 0 deletions src/macros.rs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
macro_rules! cprintln {
($surpress:expr, $stdout:expr, $fg:expr, $($rest:tt)+) => {
if !$surpress {
use std::io::Write;
use termcolor::{ColorSpec, WriteColor};

$stdout.set_color(ColorSpec::new().set_fg(Some($fg)))
.expect("Could not set the text formatting.");
writeln!($stdout, $($rest)+).expect("Could not output text.");
}
}
}

macro_rules! cprint {
($surpress:expr, $stdout:expr, $fg:expr, $($rest:tt)+) => {
if !$surpress {
use std::io::Write;
use termcolor::{ColorSpec, WriteColor};

$stdout.set_color(ColorSpec::new().set_fg(Some($fg)))
.expect("Could not set the text formatting.");
write!($stdout, $($rest)+).expect("Could not output text.");
}
}
}
245 changes: 18 additions & 227 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,23 +13,22 @@ extern crate termcolor;
extern crate generic_array;
extern crate typenum;

use clap::{Arg, ArgMatches};
use rand::OsRng;
use regex::{Regex, RegexBuilder};
use secp256k1::Secp256k1;
use std::borrow::Borrow;
use std::io::BufRead;
#[macro_use]
mod macros;
mod patterns;

use patterns::{Patterns, StringPatterns, RegexPatterns};
use std::fmt::Write;
use std::io::Write as IoWrite;
use std::sync::Mutex;
use std::thread;
use std::sync::Arc;
use std::time::Duration;
use std::fmt::Display;
use termcolor::{Color, ColorChoice, ColorSpec, WriteColor, Buffer, BufferWriter};
use generic_array::GenericArray;
use clap::{Arg, ArgMatches};
use rand::OsRng;
use regex::Regex;
use secp256k1::Secp256k1;
use termcolor::{Color, ColorChoice, Buffer, BufferWriter};
use typenum::U40;
use rayon::prelude::*;

type AddressLengthType = U40;

Expand All @@ -42,221 +41,11 @@ lazy_static! {
static ref ADDRESS_PATTERN: Regex = Regex::new(r"^[0-9a-f]{1,40}$").unwrap();
}

macro_rules! cprintln {
($surpress:expr, $stdout:expr, $fg:expr, $($rest:tt)+) => {
if !$surpress {
$stdout.set_color(ColorSpec::new().set_fg(Some($fg)))
.expect("Could not set the text formatting.");
writeln!($stdout, $($rest)+).expect("Could not output text.");
}
}
}

macro_rules! cprint {
($surpress:expr, $stdout:expr, $fg:expr, $($rest:tt)+) => {
if !$surpress {
$stdout.set_color(ColorSpec::new().set_fg(Some($fg)))
.expect("Could not set the text formatting.");
write!($stdout, $($rest)+).expect("Could not output text.");
}
}
}

struct BruteforceResult {
address: String,
private_key: String,
}

trait Pattern: Display + Send + Sync + Sized {
fn matches(&self, string: &str) -> bool;
fn parse<T: AsRef<str>>(string: T) -> Result<Self, String>;
}

impl Pattern for Regex {
fn matches(&self, string: &str) -> bool {
self.is_match(string)
}

fn parse<T: AsRef<str>>(string: T) -> Result<Self, String> {
match RegexBuilder::new(string.as_ref())
.case_insensitive(true)
.multi_line(false)
.dot_matches_new_line(false)
.ignore_whitespace(true)
.unicode(true)
.build() {
Ok(result) => return Ok(result),
Err(error) => return Err(format!("Invalid regex: {}", error)),
}
}
}

impl Pattern for String {
fn matches(&self, string: &str) -> bool {
string.starts_with(self)
}

fn parse<T: AsRef<str>>(string: T) -> Result<Self, String> {
let string = string.as_ref().to_lowercase();

if !ADDRESS_PATTERN.is_match(&string) {
return Err("Pattern contains invalid characters".to_string());
}

return Ok(string);
}
}

fn read_patterns(matches: &ArgMatches) -> Vec<String> {
if let Some(args) = matches.values_of("PATTERN") {
args.map(str::to_string).collect()
} else {
let mut result = Vec::new();
let stdin = std::io::stdin();

for line in stdin.lock().lines() {
match line {
Ok(line) => result.push(line),
Err(error) => panic!("{}", error),
}
}

result
}
}

fn parse_patterns<P: Pattern>(buffer_writer: Arc<Mutex<BufferWriter>>,
matches: &ArgMatches) -> Vec<P> {
// TODO: Use rayon (everywhere)
let mut vec: Vec<P> = Vec::new();
let raw_patterns = read_patterns(matches);

for raw_pattern in raw_patterns {
if raw_pattern.is_empty() {
continue;
}

match <P as Pattern>::parse(&raw_pattern) {
Ok(pattern) => vec.push(pattern),
Err(error) => {
let mut stdout = buffer_writer.lock().unwrap().buffer();
cprint!(matches.is_present("quiet"),
stdout,
Color::Yellow,
"Skipping pattern '{}': ",
&raw_pattern);
cprintln!(matches.is_present("quiet"),
stdout,
Color::White,
"{}",
error);
buffer_writer.lock().unwrap().print(&stdout).expect("Could not write to stdout.");
}
}
}

vec
}

trait Patterns: Sync + Send {
fn contains(&self, address: &String) -> bool;
fn len(&self) -> usize;
}

struct RegexPatterns {
vec: Vec<Regex>,
}

impl RegexPatterns {
fn new(buffer_writer: Arc<Mutex<BufferWriter>>,
matches: &ArgMatches) -> RegexPatterns {
RegexPatterns {
vec: parse_patterns(buffer_writer, matches)
}
}
}

impl Patterns for RegexPatterns {
fn contains(&self, address: &String) -> bool {
// Linear search
for pattern in &self.vec {
if pattern.matches(address) {
return true;
}
}

return false;
}

fn len(&self) -> usize {
self.vec.len()
}
}

struct StringPatterns {
// Strings of length `n` are in the `n-1`th index of this array
sorted_vecs: GenericArray<Option<Vec<String>>, AddressLengthType>,
}

impl StringPatterns {
fn new(buffer_writer: Arc<Mutex<BufferWriter>>,
matches: &ArgMatches) -> StringPatterns {
let patterns = parse_patterns::<String>(buffer_writer, matches);
// let patterns_by_len: Arc<[Mutex<Option<Vec<String>>>; ADDRESS_LENGTH]> = Arc::new([Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None)]);
let patterns_by_len: Arc<GenericArray<Mutex<Option<Vec<String>>>, AddressLengthType>> = Arc::new(arr![Mutex<Option<Vec<String>>>; Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None), Mutex::new(None)]);

patterns.par_iter()
.for_each(|pattern| {
let patterns_by_len_borrowed: &GenericArray<Mutex<Option<Vec<String>>>, AddressLengthType> = patterns_by_len.borrow();
let mut vec = patterns_by_len_borrowed[pattern.len() - 1].lock().expect("Something panicked somewhere, oops. Please report this incident to the author.");
let vec = vec.get_or_insert_with(Vec::new);

vec.push(pattern.clone());
});


let patterns_by_len_borrowed: GenericArray<Mutex<Option<Vec<String>>>, AddressLengthType> = Arc::try_unwrap(patterns_by_len).unwrap_or_else(|_| panic!("Couldn't unwrap petterns."));
let sorted_vecs = patterns_by_len_borrowed.map(|item| {
let item: Option<Vec<String>> = item.into_inner().unwrap();

item.map(|mut vec| {
vec.sort();
vec.dedup();
vec
})
});

StringPatterns {
sorted_vecs,
}
}
}

impl Patterns for StringPatterns {
fn contains(&self, address: &String) -> bool {
// Try match from shortest to longest patterns
for (index, option_vec) in self.sorted_vecs.iter().enumerate() {
if let &Some(ref vec) = option_vec {
let pattern_len = index + 1;
let target_address_slice = &address[0..pattern_len];

if vec.binary_search_by(|item| item.as_str().cmp(target_address_slice)).is_ok() {
return true;
}
}
}

return false;
}

fn len(&self) -> usize {
self.sorted_vecs.par_iter()
.filter(|opt| opt.is_some())
.map(|opt| opt.as_ref().unwrap().len())
.sum()
}
}

fn parse_color_choice(string: &str) -> Result<ColorChoice, ()> {
Ok(match string {
"always" => ColorChoice::Always,
Expand Down Expand Up @@ -328,16 +117,18 @@ patterns as regex patterns, which replaces the basic string comparison.")
let color_choice = parse_color_choice(matches.value_of("color").unwrap()).unwrap();
let buffer_writer = Arc::new(Mutex::new(BufferWriter::stdout(color_choice)));

let patterns: Arc<Patterns> = if matches.is_present("regexp") {
Arc::new(RegexPatterns::new(buffer_writer.clone(), &matches))
if matches.is_present("regexp") {
let patterns = Arc::new(RegexPatterns::new(buffer_writer.clone(), &matches));

main_pattern_type_selected(matches, quiet, buffer_writer, patterns);
} else {
Arc::new(StringPatterns::new(buffer_writer.clone(), &matches))
};
let patterns = Arc::new(StringPatterns::new(buffer_writer.clone(), &matches));

main_pattern_type_selected(matches, quiet, buffer_writer, patterns);
main_pattern_type_selected(matches, quiet, buffer_writer, patterns);
};
}

fn main_pattern_type_selected(matches: ArgMatches, quiet: bool, buffer_writer: Arc<Mutex<BufferWriter>>, patterns: Arc<Patterns>) {
fn main_pattern_type_selected<P: Patterns + 'static>(matches: ArgMatches, quiet: bool, buffer_writer: Arc<Mutex<BufferWriter>>, patterns: Arc<P>) {
if patterns.len() <= 0 {
let mut stdout = buffer_writer.lock().unwrap().buffer();
cprintln!(false,
Expand Down
73 changes: 73 additions & 0 deletions src/patterns/mod.rs
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
mod string;
mod regex;

pub use self::string::*;
pub use self::regex::*;

use std;
use std::fmt::Display;
use std::io::BufRead;
use std::sync::{Arc, Mutex};
use clap::ArgMatches;
use termcolor::{BufferWriter, Color};

trait Pattern: Display + Send + Sync + Sized {
fn matches(&self, string: &str) -> bool;
fn parse<T: AsRef<str>>(string: T) -> Result<Self, String>;
}

pub trait Patterns: Sync + Send {
fn contains(&self, address: &String) -> bool;
fn len(&self) -> usize;
}

fn read_patterns(matches: &ArgMatches) -> Vec<String> {
if let Some(args) = matches.values_of("PATTERN") {
args.map(str::to_string).collect()
} else {
let mut result = Vec::new();
let stdin = std::io::stdin();

for line in stdin.lock().lines() {
match line {
Ok(line) => result.push(line),
Err(error) => panic!("{}", error),
}
}

result
}
}

fn parse_patterns<P: Pattern>(buffer_writer: Arc<Mutex<BufferWriter>>,
matches: &ArgMatches) -> Vec<P> {
// TODO: Use rayon (everywhere)
let mut vec: Vec<P> = Vec::new();
let raw_patterns = read_patterns(matches);

for raw_pattern in raw_patterns {
if raw_pattern.is_empty() {
continue;
}

match <P as Pattern>::parse(&raw_pattern) {
Ok(pattern) => vec.push(pattern),
Err(error) => {
let mut stdout = buffer_writer.lock().unwrap().buffer();
cprint!(matches.is_present("quiet"),
stdout,
Color::Yellow,
"Skipping pattern '{}': ",
&raw_pattern);
cprintln!(matches.is_present("quiet"),
stdout,
Color::White,
"{}",
error);
buffer_writer.lock().unwrap().print(&stdout).expect("Could not write to stdout.");
}
}
}

vec
}
Loading

0 comments on commit 050f601

Please sign in to comment.