1use serde::de::{Unexpected, Visitor};
2use serde::{Deserialize, Deserializer, Serialize, Serializer};
3use sha2::{Digest, Sha256};
4use std::fmt;
5use std::fs::File;
6use std::io::{Read, Seek, SeekFrom};
7
8pub fn align(size: usize, padding: usize) -> usize {
9 ((size as usize) + padding) & !padding
10}
11
12pub fn add_padding(vec: &mut Vec<u8>, padding: usize) {
13 let real_size = vec.len();
14 vec.resize(align(real_size, padding), 0);
15}
16
17pub fn check_string_or_truncate(string: &mut String, name: &str, size: usize) {
18 if string.len() >= size {
19 println!("Warning: Truncating {} to 0x{:x}", name, size - 1);
20 string.truncate(size);
21 }
22}
23
24pub fn get_segment_data(
25 file: &mut File,
26 header: &elf::types::ProgramHeader,
27) -> std::io::Result<Vec<u8>> {
28 let mut data = vec![0; header.filesz as usize];
29 file.seek(SeekFrom::Start(header.offset))?;
30 file.read_exact(&mut data)?;
31 Ok(data)
32}
33
34pub fn compress_lz4(uncompressed_data: &mut Vec<u8>) -> std::io::Result<Vec<u8>> {
35 lz4::block::compress(&uncompressed_data[..], None, false)
36}
37
38pub fn compress_blz(uncompressed_data: &mut Vec<u8>) -> blz_nx::BlzResult<Vec<u8>> {
39 let mut compressed_data =
40 vec![0; blz_nx::get_worst_compression_buffer_size(uncompressed_data.len())];
41 let res = blz_nx::compress_raw(&mut uncompressed_data[..], &mut compressed_data[..])?;
42 compressed_data.resize(res, 0);
43 Ok(compressed_data)
44}
45
46pub fn calculate_sha256(data: &[u8]) -> std::io::Result<Vec<u8>> {
47 let mut hasher = Sha256::default();
48 hasher.update(data);
49 Ok(Vec::from(hasher.finalize().as_slice()))
50}
51
52#[derive(Default)]
53pub struct HexOrNum(pub u64);
54
55impl fmt::Debug for HexOrNum {
56 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
57 f.write_fmt(format_args!("{:#010x}", self.0))
58 }
59}
60
61impl<'de> Deserialize<'de> for HexOrNum {
62 fn deserialize<D>(deserializer: D) -> Result<HexOrNum, D::Error>
63 where
64 D: Deserializer<'de>,
65 {
66 struct HexOrNumVisitor;
67
68 impl<'a> Visitor<'a> for HexOrNumVisitor {
69 type Value = u64;
70
71 fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
72 write!(formatter, "an integer or a hex-formatted string")
73 }
74
75 fn visit_u64<E>(self, v: u64) -> Result<u64, E>
76 where
77 E: serde::de::Error,
78 {
79 Ok(v)
80 }
81
82 fn visit_str<E>(self, v: &str) -> Result<u64, E>
83 where
84 E: serde::de::Error,
85 {
86 if v.starts_with("0x") {
87 u64::from_str_radix(&v[2..], 16)
88 .map_err(|_| E::invalid_value(Unexpected::Str(v), &"a hex-encoded string"))
89 } else {
90 Err(E::invalid_value(
91 Unexpected::Str(v),
92 &"a hex-encoded string",
93 ))
94 }
95 }
96 }
97
98 let num = Deserializer::deserialize_any(deserializer, HexOrNumVisitor)?;
99 Ok(HexOrNum(num))
100 }
101}
102
103impl Serialize for HexOrNum {
104 fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
105 where
106 S: Serializer,
107 {
108 serializer.collect_str(&format_args!("{:#010x}", self.0))
109 }
110}