|
| 1 | +use anyhow::{Context, Result, bail}; |
| 2 | +use clap::{ArgAction, Parser}; |
| 3 | +use md5::{Digest, Md5}; |
| 4 | +use std::{fs, io::Write, path::PathBuf}; |
| 5 | + |
| 6 | +/// AIC firmware converter. |
| 7 | +#[derive(Parser, Debug)] |
| 8 | +#[command(author, version, about)] |
| 9 | +struct Cli { |
| 10 | + /// The input binary file. |
| 11 | + input: PathBuf, |
| 12 | + |
| 13 | + /// Output as PBP (Pre-Boot Program) format. |
| 14 | + /// |
| 15 | + /// PBP format: |
| 16 | + /// [0..=3] = b"PBP ", |
| 17 | + /// [4..=7] = checksum (u32 LE), |
| 18 | + /// [8..end] = original binary 4-byte aligned with zero paddings. |
| 19 | + #[arg(long = "pbp", action = ArgAction::SetTrue)] |
| 20 | + pbp: bool, |
| 21 | + |
| 22 | + /// Output file path. |
| 23 | + #[arg(short = 'o', long = "output")] |
| 24 | + output: PathBuf, |
| 25 | +} |
| 26 | + |
| 27 | +fn main() -> Result<()> { |
| 28 | + let cli = Cli::parse(); |
| 29 | + |
| 30 | + // Read the input binary file |
| 31 | + let bin_data = fs::read(&cli.input) |
| 32 | + .with_context(|| format!("Failed to read input file {:?}", cli.input))?; |
| 33 | + |
| 34 | + // Currently only supports -pbp mode; error if not specified |
| 35 | + if !cli.pbp { |
| 36 | + bail!("Currently only supports -pbp preprocessing, please add -pbp flag"); |
| 37 | + } |
| 38 | + |
| 39 | + let pbp_bytes = build_pbp(&bin_data)?; |
| 40 | + |
| 41 | + // Write the PBP file |
| 42 | + let mut f = fs::File::create(&cli.output) |
| 43 | + .with_context(|| format!("Failed to create output file {:?}", cli.output))?; |
| 44 | + f.write_all(&pbp_bytes) |
| 45 | + .with_context(|| format!("Failed to write output file {:?}", cli.output))?; |
| 46 | + |
| 47 | + // Pack into full image format and write .pk_pbp file |
| 48 | + let image_bytes = pack_pbp(&pbp_bytes)?; |
| 49 | + let pk_pbp_path = cli.output.with_extension("pk_pbp"); |
| 50 | + let mut f_pk = fs::File::create(&pk_pbp_path) |
| 51 | + .with_context(|| format!("Failed to create .pk_pbp file {pk_pbp_path:?}"))?; |
| 52 | + f_pk.write_all(&image_bytes) |
| 53 | + .with_context(|| format!("Failed to write .pk_pbp file {pk_pbp_path:?}"))?; |
| 54 | + |
| 55 | + Ok(()) |
| 56 | +} |
| 57 | + |
| 58 | +/// Build PBP format: |
| 59 | +/// 0..=3: "PBP " |
| 60 | +/// 4..=7: checksum (written back later) |
| 61 | +/// 8..end: bin content + zero padding to 4-byte alignment |
| 62 | +/// |
| 63 | +/// Checksum calculation: |
| 64 | +/// Sum all words of the final PBP file as u32 little-endian (mod 2^32), |
| 65 | +/// requiring the final sum == 0x0fffffff. |
| 66 | +/// First calculate partial_sum with checksum position as 0, |
| 67 | +/// then derive the checksum. |
| 68 | +fn build_pbp(bin: &[u8]) -> Result<Vec<u8>> { |
| 69 | + // 1. Align bin to 4 bytes |
| 70 | + let mut aligned = bin.to_vec(); |
| 71 | + while !aligned.len().is_multiple_of(4) { |
| 72 | + aligned.push(0); |
| 73 | + } |
| 74 | + |
| 75 | + // 2. Assemble a placeholder PBP: magic + checksum placeholder + data |
| 76 | + let total_len = aligned.len(); |
| 77 | + let mut out = Vec::with_capacity(total_len); |
| 78 | + |
| 79 | + // Data |
| 80 | + out.extend_from_slice(&aligned); |
| 81 | + |
| 82 | + // 3. Calculate the current 32-bit sum (with placeholder checksum=0) |
| 83 | + let mut sum: u64 = 0; |
| 84 | + for chunk in out.chunks_exact(4) { |
| 85 | + let word = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]); |
| 86 | + sum = (sum + word as u64) & 0xffff_ffff; |
| 87 | + } |
| 88 | + |
| 89 | + // 4. Derive the actual checksum |
| 90 | + let target: u32 = 0xffff_ffff; |
| 91 | + let checksum = target.wrapping_sub(sum as u32); |
| 92 | + |
| 93 | + // 5. Write back checksum (little-endian) |
| 94 | + out[4..8].copy_from_slice(&checksum.to_le_bytes()); |
| 95 | + |
| 96 | + // 6. Optional: Assert correctness (for debug) |
| 97 | + debug_assert!(verify_checksum(&out, target)); |
| 98 | + |
| 99 | + Ok(out) |
| 100 | +} |
| 101 | + |
| 102 | +/// Verify if the 32-bit sum of the final file is the target |
| 103 | +fn verify_checksum(buf: &[u8], target: u32) -> bool { |
| 104 | + if !buf.len().is_multiple_of(4) { |
| 105 | + return false; |
| 106 | + } |
| 107 | + let mut sum: u64 = 0; |
| 108 | + for chunk in buf.chunks_exact(4) { |
| 109 | + let word = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]); |
| 110 | + sum = (sum + word as u64) & 0xffff_ffff; |
| 111 | + } |
| 112 | + (sum as u32) == target |
| 113 | +} |
| 114 | + |
| 115 | +/// Pack PBP data into a complete Artinchip boot image format |
| 116 | +/// Format: HEAD1 + HEAD2 + DATA1 (PBP) + DATA2 (empty) + SIGN (MD5) |
| 117 | +fn pack_pbp(pbp_data: &[u8]) -> Result<Vec<u8>> { |
| 118 | + // Constants |
| 119 | + const HEAD1_SIZE: usize = 8; |
| 120 | + const HEAD2_SIZE: usize = 248; |
| 121 | + const SIGN_SIZE: usize = 16; // MD5 |
| 122 | + const ALIGNMENT: usize = 256; |
| 123 | + |
| 124 | + // Calculate sizes |
| 125 | + let pbp_total_len = pbp_data.len(); // PBP header + content |
| 126 | + let data1_len = pbp_total_len.div_ceil(ALIGNMENT) * ALIGNMENT; // Align to 256 bytes |
| 127 | + let data2_len = 0; // TODO: Implement DATA2 |
| 128 | + let signed_area_len = HEAD1_SIZE + HEAD2_SIZE + data1_len + data2_len; |
| 129 | + let total_len = signed_area_len + SIGN_SIZE; |
| 130 | + |
| 131 | + // Build HEAD1 |
| 132 | + let mut head1 = vec![0u8; HEAD1_SIZE]; |
| 133 | + head1[0..4].copy_from_slice(b"AIC "); // Magic |
| 134 | + // Checksum will be calculated later |
| 135 | + |
| 136 | + // Build HEAD2 |
| 137 | + let mut head2 = vec![0u8; HEAD2_SIZE]; |
| 138 | + // Header version: 0x00010001 |
| 139 | + head2[0..4].copy_from_slice(&0x00010001u32.to_le_bytes()); |
| 140 | + // Image length |
| 141 | + head2[4..8].copy_from_slice(&(total_len as u32).to_le_bytes()); |
| 142 | + // Firmware version: 0.0.0 (anti_rollback=0, revision=0, minor=0, major=0) |
| 143 | + // Already 0 |
| 144 | + // Loader length: 0 |
| 145 | + // Load address: 0 |
| 146 | + // Entry point: 0 |
| 147 | + // Sign algo: 0 (no signature) |
| 148 | + // Enc algo: 0 |
| 149 | + // Sign result offset |
| 150 | + head2[32..36].copy_from_slice(&(signed_area_len as u32).to_le_bytes()); |
| 151 | + // Sign result length: 16 |
| 152 | + head2[36..40].copy_from_slice(&16u32.to_le_bytes()); |
| 153 | + // Other offsets: 0 |
| 154 | + // PBP offset: 256 |
| 155 | + head2[64..68].copy_from_slice(&256u32.to_le_bytes()); |
| 156 | + // PBP length: pbp_data.len() |
| 157 | + head2[68..72].copy_from_slice(&(pbp_data.len() as u32).to_le_bytes()); |
| 158 | + // Padding: already 0 |
| 159 | + |
| 160 | + // Build DATA1: PBP + padding |
| 161 | + let mut data1 = vec![0u8; data1_len]; |
| 162 | + data1[0..pbp_data.len()].copy_from_slice(pbp_data); |
| 163 | + |
| 164 | + // TODO DATA2: empty |
| 165 | + |
| 166 | + // SIGN: MD5 of HEAD2 + DATA1 + DATA2 |
| 167 | + let signed_data = [head2.as_slice(), data1.as_slice()].concat(); |
| 168 | + let mut hasher = Md5::new(); |
| 169 | + hasher.update(&signed_data); |
| 170 | + let sign = hasher.finalize().to_vec(); |
| 171 | + |
| 172 | + let signed_area = [head1.as_slice(), head2.as_slice(), data1.as_slice()].concat(); |
| 173 | + |
| 174 | + let mut result = vec![0u8; total_len]; |
| 175 | + result[0..signed_area_len].copy_from_slice(&signed_area); |
| 176 | + result[signed_area_len..].copy_from_slice(&sign); |
| 177 | + |
| 178 | + // Calculate HEAD1 checksum |
| 179 | + let mut sum: u64 = 0; |
| 180 | + for chunk in result.chunks_exact(4) { |
| 181 | + let word = u32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]); |
| 182 | + sum = (sum + word as u64) & 0xffff_ffff; |
| 183 | + } |
| 184 | + |
| 185 | + let target: u32 = 0xffff_ffff; |
| 186 | + let head1_checksum = target.wrapping_sub(sum as u32); |
| 187 | + |
| 188 | + result[4..8].copy_from_slice(&head1_checksum.to_le_bytes()); |
| 189 | + |
| 190 | + Ok(result) |
| 191 | +} |
0 commit comments