1+ use pyo3:: prelude:: * ;
2+ use pyo3:: exceptions:: PyIOError ;
3+ use rayon:: prelude:: * ;
4+ use std:: fs:: File ;
5+ use std:: io:: { BufRead , BufReader , Read } ;
6+ use std:: sync:: atomic:: { AtomicBool , AtomicU64 , Ordering } ;
7+ use std:: sync:: Arc ;
8+ use std:: time:: { Duration , Instant } ;
9+ use zip:: ZipArchive ;
10+
11+ #[ pyclass]
12+ pub struct BruteForce {
13+ actual_trys : Arc < AtomicU64 > ,
14+ found : Arc < AtomicBool > ,
15+ result : Arc < parking_lot:: Mutex < Option < String > > > ,
16+ }
17+
18+ #[ pymethods]
19+ impl BruteForce {
20+ #[ new]
21+ #[ pyo3( signature = ( src=None ) ) ]
22+ fn new ( src : Option < String > ) -> Self {
23+ let _ = src;
24+ BruteForce {
25+ actual_trys : Arc :: new ( AtomicU64 :: new ( 0 ) ) ,
26+ found : Arc :: new ( AtomicBool :: new ( false ) ) ,
27+ result : Arc :: new ( parking_lot:: Mutex :: new ( None ) ) ,
28+ }
29+ }
30+
31+ fn get_tries ( & self ) -> u64 {
32+ self . actual_trys . load ( Ordering :: Relaxed )
33+ }
34+
35+ fn is_found ( & self ) -> bool {
36+ self . found . load ( Ordering :: Relaxed )
37+ }
38+
39+ fn get_result ( & self ) -> Option < String > {
40+ self . result . lock ( ) . clone ( )
41+ }
42+
43+ fn reset ( & self ) {
44+ self . actual_trys . store ( 0 , Ordering :: Relaxed ) ;
45+ self . found . store ( false , Ordering :: Relaxed ) ;
46+ * self . result . lock ( ) = None ;
47+ }
48+
49+ #[ pyo3( signature = ( zip_path, wordlist_path, callback=None , batch_size=None ) ) ]
50+ fn crack (
51+ & self ,
52+ py : Python ,
53+ zip_path : String ,
54+ wordlist_path : String ,
55+ callback : Option < PyObject > ,
56+ batch_size : Option < usize > ,
57+ ) -> PyResult < Option < String > > {
58+ self . reset ( ) ;
59+
60+ let batch_size = batch_size. unwrap_or ( 50000 ) ;
61+
62+ let actual_trys = self . actual_trys . clone ( ) ;
63+ let found = self . found . clone ( ) ;
64+ let result = self . result . clone ( ) ;
65+
66+ let final_result = py. allow_threads ( || {
67+ let file = File :: open ( & wordlist_path)
68+ . map_err ( |e| format ! ( "Error opening wordlist: {}" , e) ) ?;
69+
70+ let reader = BufReader :: with_capacity ( 128 * 1024 , file) ;
71+ let mut passwords = Vec :: with_capacity ( batch_size) ;
72+
73+ let mut batches_processed = 0u64 ;
74+ let start_time = Instant :: now ( ) ;
75+ let mut last_callback_time = Instant :: now ( ) ;
76+
77+ for line in reader. lines ( ) {
78+ if found. load ( Ordering :: Relaxed ) {
79+ break ;
80+ }
81+
82+ let password = line. map_err ( |e| format ! ( "Error reading line: {}" , e) ) ?;
83+ passwords. push ( password) ;
84+
85+ if passwords. len ( ) >= batch_size {
86+ let result_found = Self :: process_batch_static (
87+ & zip_path,
88+ & passwords,
89+ & actual_trys,
90+ & found,
91+ & result,
92+ ) ?;
93+
94+ batches_processed += 1 ;
95+
96+ if result_found. is_some ( ) {
97+ return Ok ( result_found) ;
98+ }
99+
100+ passwords. clear ( ) ;
101+
102+ if batches_processed % 5 == 0 && last_callback_time. elapsed ( ) >= Duration :: from_millis ( 200 ) {
103+ std:: thread:: sleep ( Duration :: from_micros ( 100 ) ) ;
104+ last_callback_time = Instant :: now ( ) ;
105+ }
106+ }
107+ }
108+
109+ if !passwords. is_empty ( ) && !found. load ( Ordering :: Relaxed ) {
110+ let result_found = Self :: process_batch_static (
111+ & zip_path,
112+ & passwords,
113+ & actual_trys,
114+ & found,
115+ & result,
116+ ) ?;
117+
118+ if result_found. is_some ( ) {
119+ return Ok ( result_found) ;
120+ }
121+ }
122+
123+ Ok ( None )
124+ } ) . map_err ( |e : String | PyIOError :: new_err ( e) ) ?;
125+
126+ if let Some ( ref cb) = callback {
127+ let tries = actual_trys. load ( Ordering :: Relaxed ) ;
128+ cb. call1 ( py, ( tries, ) ) ?;
129+ }
130+
131+ Ok ( final_result)
132+ }
133+
134+ #[ getter]
135+ fn actual_trys ( & self ) -> u64 {
136+ self . get_tries ( )
137+ }
138+ }
139+
140+ impl BruteForce {
141+ fn process_batch_static (
142+ zip_path : & str ,
143+ passwords : & [ String ] ,
144+ actual_trys : & Arc < AtomicU64 > ,
145+ found : & Arc < AtomicBool > ,
146+ result : & Arc < parking_lot:: Mutex < Option < String > > > ,
147+ ) -> Result < Option < String > , String > {
148+ let found_password = passwords. par_iter ( ) . find_map_any ( |password| {
149+ if found. load ( Ordering :: Relaxed ) {
150+ return None ;
151+ }
152+
153+ actual_trys. fetch_add ( 1 , Ordering :: Relaxed ) ;
154+
155+ let file = File :: open ( zip_path) . ok ( ) ?;
156+ let mut archive = ZipArchive :: new ( file) . ok ( ) ?;
157+
158+ if archive. len ( ) == 0 {
159+ return None ;
160+ }
161+
162+ let entry_result = archive. by_index_decrypt ( 0 , password. as_bytes ( ) ) . ok ( ) ?;
163+ let mut entry = entry_result. ok ( ) ?;
164+ let mut buffer = Vec :: new ( ) ;
165+
166+ if entry. read_to_end ( & mut buffer) . is_ok ( ) {
167+ Some ( password. clone ( ) )
168+ } else {
169+ None
170+ }
171+
172+ } ) ;
173+
174+ if let Some ( ref password) = found_password {
175+ found. store ( true , Ordering :: Relaxed ) ;
176+ * result. lock ( ) = Some ( password. clone ( ) ) ;
177+ }
178+
179+ Ok ( found_password)
180+ }
181+ }
182+
183+ #[ pymodule]
184+ fn fast_zip ( _py : Python < ' _ > , m : & PyModule ) -> PyResult < ( ) > {
185+ m. add_class :: < BruteForce > ( ) ?;
186+ m. add_function ( wrap_pyfunction ! ( bruteforce_zip, m) ?) ?;
187+ Ok ( ( ) )
188+ }
0 commit comments