1use aes::cipher::consts::U16;
2use aes::cipher::generic_array::GenericArray;
3use chrono::Local;
4use log::{info, warn, LevelFilter};
5use std::fs::File;
6use std::{fs, io};
7use std::io::{ Read, Seek, SeekFrom};
8use log4rs::{
9 append::{console::ConsoleAppender, file::FileAppender},
10 config::{Appender, Config, Root},
11 encode::pattern::PatternEncoder,
12};
13
14
15#[cfg(unix)]
16use std::os::unix::fs::FileExt as _;
17#[cfg(windows)]
18use std::os::windows::fs::FileExt as _;
19use std::path::Path;
20
21
22#[cfg(unix)]
23pub fn read_exact_at(file: &File, mut buf: &mut [u8], mut off: u64) -> io::Result<()> {
24 while !buf.is_empty() {
25 let n = file.read_at(buf, off)?;
26 if n == 0 { return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "read_at=0")); }
27 buf = &mut buf[n..];
28 off += n as u64;
29 }
30 Ok(())
31}
32#[cfg(windows)]
33pub fn read_exact_at(file: &File, mut buf: &mut [u8], mut off: u64) -> io::Result<()> {
34 while !buf.is_empty() {
35 let n = file.seek_read(buf, off)?;
36 if n == 0 { return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "seek_read=0")); }
37 buf = &mut buf[n..];
38 off += n as u64;
39 }
40 Ok(())
41}
42
43#[cfg(unix)]
44pub fn write_all_at(file: &File, mut buf: &[u8], mut off: u64) -> io::Result<()> {
45 while !buf.is_empty() {
46 let n = file.write_at(buf, off)?;
47 if n == 0 { return Err(io::Error::new(io::ErrorKind::WriteZero, "write_at=0")); }
48 buf = &buf[n..];
49 off += n as u64;
50 }
51 Ok(())
52}
53#[cfg(windows)]
54pub fn write_all_at(file: &File, mut buf: &[u8], mut off: u64) -> io::Result<()> {
55 while !buf.is_empty() {
56 let n = file.seek_write(buf, off)?;
57 if n == 0 { return Err(io::Error::new(io::ErrorKind::WriteZero, "seek_write=0")); }
58 buf = &buf[n..];
59 off += n as u64;
60 }
61 Ok(())
62}
63
64
65pub struct Region {
66 start: u64,
67 end: u64,
68}
69
70pub fn key_validation(key: &str) -> bool {
72 let stripped_key: String = key.chars().filter(|c| !c.is_whitespace()).collect();
73
74 if stripped_key.len() != 32 {
75 println!("Invalid key length: {}", stripped_key.len());
76 warn!("Invalid key length: {}", stripped_key.len());
77 return false;
78 }
79
80 if !stripped_key.chars().all(|c| c.is_ascii_hexdigit()) {
81 println!("Key contains invalid characters");
82 warn!("Key contains invalid characters");
83 return false;
84 }
85
86 info!("Key is valid");
87 true
88}
89
90#[inline(always)]
92pub fn generate_iv(sector: u64) -> GenericArray<u8, U16> {
93 let mut iv_bytes = [0u8; 16];
94 iv_bytes[12] = (sector >> 24) as u8;
95 iv_bytes[13] = (sector >> 16) as u8;
96 iv_bytes[14] = (sector >> 8) as u8;
97 iv_bytes[15] = sector as u8;
98 GenericArray::clone_from_slice(&iv_bytes)
99}
100#[inline(always)]
101pub fn is_encrypted(regions: &[Region], sector: u64, sector_data: &[u8]) -> bool {
102 if sector_data.iter().all(|&b| b == 0) {
103 return false;
104 }
105 regions.iter().any(|r| sector >= r.start && sector < r.end)
106}
107
108
109
110pub fn extract_regions<R: Read + Seek>(reader: &mut R) -> io::Result<Vec<Region>> {
112 let mut header = [0u8; 4096];
113 reader.seek(SeekFrom::Start(0))?;
114 reader.read_exact(&mut header)?;
115 let num_normal_regions = u32::from_be_bytes(header[0..4].try_into().unwrap()) as usize;
116 let regions_count = (num_normal_regions * 2) - 1;
117 let mut regions = Vec::with_capacity(regions_count);
118
119 let mut is_encrypted = false;
120 for i in 0..regions_count {
121 let region_offset = 4 + i * 8;
122 let start_sector =
123 u32::from_be_bytes(header[region_offset..region_offset + 4].try_into().unwrap());
124 let end_sector = u32::from_be_bytes(
125 header[region_offset + 4..region_offset + 8]
126 .try_into()
127 .unwrap(),
128 );
129
130 regions.push(Region {
131 start: start_sector as u64,
132 end: end_sector as u64,
133 });
134
135 is_encrypted = !is_encrypted;
136 }
137
138 Ok(regions)
139}
140
141
142pub fn setup_logging() -> Result<(), Box<dyn std::error::Error>> {
143 let log_dir = Path::new("log");
144 fs::create_dir_all(log_dir)?;
145 let now = Local::now();
146 let log_file_name = format!("log/{}.log", now.format("%Y-%m-%d_%H-%M-%S"));
147
148 let fmt = "{d(%Y-%m-%d %H:%M:%S)} [{l}] - {m}\n";
149
150 let stdout = ConsoleAppender::builder()
151 .encoder(Box::new(PatternEncoder::new(fmt)))
152 .build();
153
154 let logfile = FileAppender::builder()
155 .encoder(Box::new(PatternEncoder::new(fmt)))
156 .build(log_file_name)?;
157
158 let config = Config::builder()
159 .appender(Appender::builder().build("stdout", Box::new(stdout)))
160 .appender(Appender::builder().build("logfile", Box::new(logfile)))
161 .build(
162 Root::builder()
163 .appender("stdout")
164 .appender("logfile")
165 .build(LevelFilter::Trace),
166 )?;
167
168 log4rs::init_config(config)?;
169 Ok(())
170}