207 lines
6.0 kB
1
use anyhow::Result;
2
use log::{debug, info};
3
use rayon::prelude::*;
4
use std::sync::{Arc, Mutex};
5
6
use crate::memory::MemoryReader;
7
use crate::process::{MemoryRegion, Process};
8
9
#[derive(Clone, Debug)]
10
pub enum ScanValueType {
11
#[allow(dead_code)]
12
U8(u8),
13
#[allow(dead_code)]
14
U16(u16),
15
#[allow(dead_code)]
16
U32(u32),
17
#[allow(dead_code)]
18
U64(u64),
19
#[allow(dead_code)]
20
I8(i8),
21
#[allow(dead_code)]
22
I16(i16),
23
I32(i32),
24
#[allow(dead_code)]
25
I64(i64),
26
#[allow(dead_code)]
27
F32(f32),
28
#[allow(dead_code)]
29
F64(f64),
30
}
31
32
#[derive(Clone, Debug)]
33
pub struct MemoryMatch {
34
pub address: usize,
35
#[allow(dead_code)]
36
pub value_type: ScanValueType,
37
}
38
39
pub struct MemoryScanner {
40
reader: MemoryReader,
41
regions: Vec<MemoryRegion>,
42
matches: Vec<MemoryMatch>,
43
}
44
45
impl MemoryScanner {
46
pub fn new(process: Process) -> Result<Self> {
47
let regions = process.get_memory_maps()?;
48
let reader = MemoryReader::new(process);
49
50
println!(
51
"Successfully attached to process. Found {} readable memory regions.",
52
regions.len()
53
);
54
55
Ok(Self {
56
reader,
57
regions,
58
matches: Vec::new(),
59
})
60
}
61
62
pub fn scan_for_value(&mut self, value: ScanValueType) -> Result<()> {
63
let matches = Arc::new(Mutex::new(Vec::new()));
64
65
self.regions
66
.par_iter()
67
.enumerate()
68
.for_each(|(idx, region)| {
69
println!(
70
"{:02}/{:03} searching {:x} - {:x}...",
71
idx + 1,
72
self.regions.len(),
73
region.start_addr,
74
region.end_addr
75
);
76
77
match self.scan_region(region, &value) {
78
Ok(region_matches) => {
79
let mut matches_guard = matches.lock().unwrap();
80
matches_guard.extend(region_matches);
81
}
82
Err(e) => {
83
debug!(
84
"Error scanning region {:x}-{:x}: {}",
85
region.start_addr, region.end_addr, e
86
);
87
}
88
}
89
90
print!("........ok\n");
91
});
92
93
self.matches = Arc::try_unwrap(matches)
94
.expect("Failed to unwrap Arc")
95
.into_inner()
96
.expect("Failed to unlock Mutex");
97
98
info!("we currently have {} matches.", self.matches.len());
99
Ok(())
100
}
101
102
pub fn get_matches_count(&self) -> usize {
103
self.matches.len()
104
}
105
106
fn scan_region(
107
&self,
108
region: &MemoryRegion,
109
value: &ScanValueType,
110
) -> Result<Vec<MemoryMatch>> {
111
// Basic implementation that scans a single region for matches
112
let mut region_matches = Vec::new();
113
114
match self.reader.read_region(region) {
115
Ok(memory) => {
116
// Simple implementation searching for i32 values
117
if let ScanValueType::I32(target) = value {
118
let target_bytes = target.to_ne_bytes();
119
120
for i in (0..memory.len().saturating_sub(4)).step_by(4) {
121
let mut matches = true;
122
for j in 0..4 {
123
if memory[i + j] != target_bytes[j] {
124
matches = false;
125
break;
126
}
127
}
128
129
if matches {
130
region_matches.push(MemoryMatch {
131
address: region.start_addr + i,
132
value_type: value.clone(),
133
});
134
}
135
}
136
}
137
// TODO missing handlers for all other types
138
}
139
Err(e) => {
140
println!(
141
"Error reading region {:x}-{:x}: {}",
142
region.start_addr, region.end_addr, e
143
);
144
}
145
}
146
147
Ok(region_matches)
148
}
149
pub fn filter_matches(&mut self, value: ScanValueType) -> Result<()> {
150
let old_matches = std::mem::take(&mut self.matches);
151
152
for m in old_matches {
153
let addr = m.address;
154
155
let data = match &value {
156
ScanValueType::I32(_) => {
157
match self.reader.read_memory(addr, 4) {
158
Ok(data) => data,
159
Err(_) => continue, // Skip if we can't read
160
}
161
}
162
_ => continue,
163
};
164
165
let matched = match &value {
166
ScanValueType::I32(target) => {
167
if data.len() >= 4 {
168
let value_bytes = target.to_ne_bytes();
169
data[0] == value_bytes[0]
170
&& data[1] == value_bytes[1]
171
&& data[2] == value_bytes[2]
172
&& data[3] == value_bytes[3]
173
} else {
174
false
175
}
176
}
177
_ => false,
178
};
179
180
if matched {
181
self.matches.push(MemoryMatch {
182
address: addr,
183
value_type: value.clone(),
184
});
185
}
186
}
187
188
println!("..........ok");
189
info!("we currently have {} matches.", self.matches.len());
190
Ok(())
191
}
192
193
pub fn set_value(&self, index: usize, value: i32) -> Result<()> {
194
if index >= self.matches.len() {
195
anyhow::bail!(
196
"Index {} out of range (have {} matches)",
197
index,
198
self.matches.len()
199
);
200
}
201
202
let address = self.matches[index].address;
203
let data = value.to_ne_bytes();
204
205
self.reader.write_memory(address, &data)
206
}
207
}
208