use std::{borrow::Cow, io, mem::replace, ops::Range, str::Utf8Error};
use serde::{Deserialize, Serialize};
use snafu::{Backtrace, Snafu};
use super::{
raw::{
AutoloadInfo, AutoloadInfoEntry, AutoloadKind, BuildInfo, HmacSha1Signature, HmacSha1SignatureError,
RawAutoloadInfoError, RawBuildInfoError, NITROCODE_BYTES,
},
Autoload, OverlayTable,
};
use crate::{
compress::lz77::{Lz77, Lz77DecompressError},
crc::CRC_16_MODBUS,
crypto::blowfish::{Blowfish, BlowfishError, BlowfishKey, BlowfishLevel},
rom::LibraryEntry,
};
#[derive(Clone)]
pub struct Arm9<'a> {
data: Cow<'a, [u8]>,
offsets: Arm9Offsets,
originally_compressed: bool,
originally_encrypted: bool,
}
#[derive(Serialize, Deserialize, Clone, Copy)]
pub struct Arm9Offsets {
pub base_address: u32,
pub entry_function: u32,
pub build_info: u32,
pub autoload_callback: u32,
pub overlay_signatures: u32,
}
const SECURE_AREA_ID: [u8; 8] = [0xff, 0xde, 0xff, 0xe7, 0xff, 0xde, 0xff, 0xe7];
const SECURE_AREA_ENCRY_OBJ: &[u8] = "encryObj".as_bytes();
const LZ77: Lz77 = Lz77 {};
const COMPRESSION_START: usize = 0x4000;
#[derive(Debug, Snafu)]
pub enum Arm9Error {
#[snafu(display("expected {expected:#x} bytes for secure area but had only {actual:#x}:\n{backtrace}"))]
DataTooSmall {
expected: usize,
actual: usize,
backtrace: Backtrace,
},
#[snafu(transparent)]
Blowfish {
source: BlowfishError,
},
#[snafu(display("invalid encryption, 'encryObj' not found"))]
NotEncryObj {
backtrace: Backtrace,
},
#[snafu(transparent)]
RawBuildInfo {
source: RawBuildInfoError,
},
#[snafu(transparent)]
Lz77Decompress {
source: Lz77DecompressError,
},
#[snafu(transparent)]
Io {
source: io::Error,
},
}
#[derive(Debug, Snafu)]
pub enum Arm9AutoloadError {
#[snafu(transparent)]
RawBuildInfo {
source: RawBuildInfoError,
},
#[snafu(transparent)]
RawAutoloadInfo {
source: RawAutoloadInfoError,
},
#[snafu(display("ARM9 program must be decompressed before accessing autoload blocks:\n{backtrace}"))]
Compressed {
backtrace: Backtrace,
},
#[snafu(display("autoload block {kind} could not be found:\n{backtrace}"))]
NotFound {
kind: AutoloadKind,
backtrace: Backtrace,
},
}
#[derive(Debug, Snafu)]
pub enum Arm9OverlaySignaturesError {
#[snafu(transparent)]
HmacSha1Signature {
source: HmacSha1SignatureError,
},
#[snafu(transparent)]
RawBuildInfo {
source: RawBuildInfoError,
},
#[snafu(display("ARM9 program must be decompressed before accessing overlay signatures:\n{backtrace}"))]
OverlaySignaturesCompressed {
backtrace: Backtrace,
},
}
#[derive(Debug, Snafu)]
pub enum Arm9HmacSha1KeyError {
#[snafu(transparent)]
RawBuildInfo {
source: RawBuildInfoError,
},
#[snafu(display("ARM9 program must be decompressed before accessing HMAC-SHA1 key:\n{backtrace}"))]
HmacSha1KeyCompressed {
backtrace: Backtrace,
},
}
pub struct Arm9WithTcmsOptions {
pub originally_compressed: bool,
pub originally_encrypted: bool,
}
impl<'a> Arm9<'a> {
pub fn new<T: Into<Cow<'a, [u8]>>>(data: T, offsets: Arm9Offsets) -> Result<Self, RawBuildInfoError> {
let mut arm9 = Arm9 { data: data.into(), offsets, originally_compressed: false, originally_encrypted: false };
arm9.originally_compressed = arm9.is_compressed()?;
arm9.originally_encrypted = arm9.is_encrypted();
Ok(arm9)
}
pub fn with_two_tcms(
mut data: Vec<u8>,
itcm: Autoload,
dtcm: Autoload,
offsets: Arm9Offsets,
options: Arm9WithTcmsOptions,
) -> Result<Self, RawBuildInfoError> {
let autoload_infos = [*itcm.info().entry(), *dtcm.info().entry()];
let autoload_blocks = data.len() as u32 + offsets.base_address;
data.extend(itcm.into_data().iter());
data.extend(dtcm.into_data().iter());
let autoload_infos_start = data.len() as u32 + offsets.base_address;
data.extend(bytemuck::bytes_of(&autoload_infos));
let autoload_infos_end = data.len() as u32 + offsets.base_address;
let Arm9WithTcmsOptions { originally_compressed, originally_encrypted } = options;
let mut arm9 = Self { data: data.into(), offsets, originally_compressed, originally_encrypted };
let build_info = arm9.build_info_mut()?;
build_info.autoload_blocks = autoload_blocks;
build_info.autoload_infos_start = autoload_infos_start;
build_info.autoload_infos_end = autoload_infos_end;
Ok(arm9)
}
pub fn with_autoloads(
mut data: Vec<u8>,
autoloads: &[Autoload],
offsets: Arm9Offsets,
options: Arm9WithTcmsOptions,
) -> Result<Self, RawBuildInfoError> {
let autoload_blocks = data.len() as u32 + offsets.base_address;
for autoload in autoloads {
data.extend(autoload.full_data());
}
let autoload_infos_start = data.len() as u32 + offsets.base_address;
for autoload in autoloads {
data.extend(bytemuck::bytes_of(autoload.info().entry()));
}
let autoload_infos_end = data.len() as u32 + offsets.base_address;
let Arm9WithTcmsOptions { originally_compressed, originally_encrypted } = options;
let mut arm9 = Self { data: data.into(), offsets, originally_compressed, originally_encrypted };
let build_info = arm9.build_info_mut()?;
build_info.autoload_blocks = autoload_blocks;
build_info.autoload_infos_start = autoload_infos_start;
build_info.autoload_infos_end = autoload_infos_end;
Ok(arm9)
}
pub fn is_encrypted(&self) -> bool {
self.data.len() < 8 || self.data[0..8] != SECURE_AREA_ID
}
pub fn decrypt(&mut self, key: &BlowfishKey, gamecode: u32) -> Result<(), Arm9Error> {
if !self.is_encrypted() {
return Ok(());
}
if self.data.len() < 0x4000 {
DataTooSmallSnafu { expected: 0x4000usize, actual: self.data.len() }.fail()?;
}
let mut secure_area = [0u8; 0x4000];
secure_area.clone_from_slice(&self.data[0..0x4000]);
let blowfish = Blowfish::new(key, gamecode, BlowfishLevel::Level2);
blowfish.decrypt(&mut secure_area[0..8])?;
let blowfish = Blowfish::new(key, gamecode, BlowfishLevel::Level3);
blowfish.decrypt(&mut secure_area[0..0x800])?;
if &secure_area[0..8] != SECURE_AREA_ENCRY_OBJ {
NotEncryObjSnafu {}.fail()?;
}
secure_area[0..8].copy_from_slice(&SECURE_AREA_ID);
self.data.to_mut()[0..0x4000].copy_from_slice(&secure_area);
Ok(())
}
pub fn encrypt(&mut self, key: &BlowfishKey, gamecode: u32) -> Result<(), Arm9Error> {
if self.is_encrypted() {
return Ok(());
}
if self.data.len() < 0x4000 {
DataTooSmallSnafu { expected: 0x4000usize, actual: self.data.len() }.fail()?;
}
if self.data[0..8] != SECURE_AREA_ID {
NotEncryObjSnafu {}.fail()?;
}
let secure_area = self.encrypted_secure_area(key, gamecode);
self.data.to_mut()[0..0x4000].copy_from_slice(&secure_area);
Ok(())
}
pub fn encrypted_secure_area(&self, key: &BlowfishKey, gamecode: u32) -> [u8; 0x4000] {
let mut secure_area = [0u8; 0x4000];
secure_area.copy_from_slice(&self.data[0..0x4000]);
if self.is_encrypted() {
return secure_area;
}
secure_area[0..8].copy_from_slice(SECURE_AREA_ENCRY_OBJ);
let blowfish = Blowfish::new(key, gamecode, BlowfishLevel::Level3);
blowfish.encrypt(&mut secure_area[0..0x800]).unwrap();
let blowfish = Blowfish::new(key, gamecode, BlowfishLevel::Level2);
blowfish.encrypt(&mut secure_area[0..8]).unwrap();
secure_area
}
pub fn secure_area_crc(&self, key: &BlowfishKey, gamecode: u32) -> u16 {
let secure_area = self.encrypted_secure_area(key, gamecode);
CRC_16_MODBUS.checksum(&secure_area)
}
pub fn build_info(&self) -> Result<&BuildInfo, RawBuildInfoError> {
BuildInfo::borrow_from_slice(&self.data[self.offsets.build_info as usize..])
}
pub fn build_info_mut(&mut self) -> Result<&mut BuildInfo, RawBuildInfoError> {
BuildInfo::borrow_from_slice_mut(&mut self.data.to_mut()[self.offsets.build_info as usize..])
}
pub fn libraries(&self) -> Result<Box<[LibraryEntry<'_>]>, Utf8Error> {
let libraries_start = self.offsets.build_info as usize + size_of::<BuildInfo>();
let mut address = self.base_address() + libraries_start as u32;
let mut data = &self.data[libraries_start..];
let mut libraries = Vec::new();
'outer: while data[0] == b'[' {
let Some((end_pos, _)) = data.iter().enumerate().find(|(_, c)| **c == b']') else {
break;
};
let version_string = str::from_utf8(&data[..end_pos + 1])?;
libraries.push(LibraryEntry::new(address, version_string));
address += end_pos as u32 + 1;
data = &data[end_pos + 1..];
loop {
match data[0] {
b'\0' => {
address += 1;
data = &data[1..]
}
b'[' => break,
_ => break 'outer,
}
}
}
Ok(libraries.into_boxed_slice())
}
pub fn is_compressed(&self) -> Result<bool, RawBuildInfoError> {
Ok(self.build_info()?.is_compressed())
}
pub fn decompress(&mut self) -> Result<(), Arm9Error> {
if !self.is_compressed()? {
return Ok(());
}
let data: Cow<[u8]> = LZ77.decompress(&self.data)?.into_vec().into();
let old_data = replace(&mut self.data, data);
let build_info = match self.build_info_mut() {
Ok(build_info) => build_info,
Err(e) => {
self.data = old_data;
return Err(e.into());
}
};
build_info.compressed_code_end = 0;
Ok(())
}
pub fn compress(&mut self) -> Result<(), Arm9Error> {
if self.is_compressed()? {
return Ok(());
}
let data: Cow<[u8]> = LZ77.compress(&self.data, COMPRESSION_START)?.into_vec().into();
let length = data.len();
let old_data = replace(&mut self.data, data);
let base_address = self.base_address();
let build_info = match self.build_info_mut() {
Ok(build_info) => build_info,
Err(e) => {
self.data = old_data;
return Err(e.into());
}
};
build_info.compressed_code_end = base_address + length as u32;
Ok(())
}
fn get_autoload_info_entries(&self, build_info: &BuildInfo) -> Result<&[AutoloadInfoEntry], Arm9AutoloadError> {
let start = (build_info.autoload_infos_start - self.base_address()) as usize;
let end = (build_info.autoload_infos_end - self.base_address()) as usize;
let autoload_info = AutoloadInfoEntry::borrow_from_slice(&self.data[start..end])?;
Ok(autoload_info)
}
pub fn autoload_infos(&self) -> Result<Vec<AutoloadInfo>, Arm9AutoloadError> {
let build_info: &BuildInfo = self.build_info()?;
if build_info.is_compressed() {
CompressedSnafu {}.fail()?;
}
Ok(self
.get_autoload_info_entries(build_info)?
.iter()
.enumerate()
.map(|(index, entry)| AutoloadInfo::new(*entry, index as u32))
.collect())
}
pub fn autoloads(&self) -> Result<Box<[Autoload<'_>]>, Arm9AutoloadError> {
let build_info = self.build_info()?;
if build_info.is_compressed() {
CompressedSnafu {}.fail()?;
}
let autoload_infos = self.autoload_infos()?;
let mut autoloads = vec![];
let mut load_offset = build_info.autoload_blocks - self.base_address();
for autoload_info in autoload_infos {
let start = load_offset as usize;
let end = start + autoload_info.code_size() as usize;
let data = &self.data[start..end];
autoloads.push(Autoload::new(data, autoload_info));
load_offset += autoload_info.code_size();
}
Ok(autoloads.into_boxed_slice())
}
pub fn num_unknown_autoloads(&self) -> Result<usize, Arm9AutoloadError> {
Ok(self.autoloads()?.iter().filter(|a| matches!(a.kind(), AutoloadKind::Unknown(_))).count())
}
pub fn hmac_sha1_key(&self) -> Result<Option<[u8; 64]>, Arm9HmacSha1KeyError> {
if self.is_compressed()? {
HmacSha1KeyCompressedSnafu {}.fail()?
}
let Some((i, _)) = self.data.chunks(4).enumerate().filter(|(_, chunk)| *chunk == NITROCODE_BYTES).nth(1) else {
return Ok(None);
};
let start = i * 4;
let end = start + 64;
if end > self.data.len() {
return Ok(None);
}
let mut key = [0u8; 64];
key.copy_from_slice(&self.data[start..end]);
Ok(Some(key))
}
fn overlay_table_signature_range(&self) -> Result<Option<Range<usize>>, Arm9OverlaySignaturesError> {
let overlay_signatures_offset = self.overlay_signatures_offset() as usize;
if overlay_signatures_offset == 0 {
return Ok(None);
}
if self.is_compressed()? {
OverlaySignaturesCompressedSnafu {}.fail()?;
}
let start = overlay_signatures_offset - size_of::<HmacSha1Signature>();
let end = overlay_signatures_offset;
if end > self.data.len() {
return Ok(None);
}
Ok(Some(start..end))
}
pub fn overlay_table_signature(&self) -> Result<Option<&HmacSha1Signature>, Arm9OverlaySignaturesError> {
let Some(range) = self.overlay_table_signature_range()? else {
return Ok(None);
};
let data = &self.data[range];
let signature = HmacSha1Signature::borrow_from_slice(data)?;
Ok(Some(signature.first().unwrap()))
}
pub fn overlay_table_signature_mut(&mut self) -> Result<Option<&mut HmacSha1Signature>, Arm9OverlaySignaturesError> {
let Some(range) = self.overlay_table_signature_range()? else {
return Ok(None);
};
let data = &mut self.data.to_mut()[range];
let signature = HmacSha1Signature::borrow_from_slice_mut(data)?;
Ok(Some(signature.first_mut().unwrap()))
}
fn overlay_signatures_range(&self, num_overlays: usize) -> Result<Option<Range<usize>>, Arm9OverlaySignaturesError> {
let start = self.overlay_signatures_offset() as usize;
if start == 0 {
return Ok(None);
}
if self.is_compressed()? {
OverlaySignaturesCompressedSnafu {}.fail()?;
}
let end = start + size_of::<HmacSha1Signature>() * num_overlays;
if end > self.data.len() {
return Ok(None);
}
Ok(Some(start..end))
}
pub fn overlay_signatures(&self, num_overlays: usize) -> Result<Option<&[HmacSha1Signature]>, Arm9OverlaySignaturesError> {
let Some(range) = self.overlay_signatures_range(num_overlays)? else {
return Ok(None);
};
let data = &self.data[range];
Ok(Some(HmacSha1Signature::borrow_from_slice(data)?))
}
pub fn overlay_signatures_mut(
&mut self,
num_overlays: usize,
) -> Result<Option<&mut [HmacSha1Signature]>, Arm9OverlaySignaturesError> {
let Some(range) = self.overlay_signatures_range(num_overlays)? else {
return Ok(None);
};
let data = &mut self.data.to_mut()[range];
Ok(Some(HmacSha1Signature::borrow_from_slice_mut(data)?))
}
pub fn code(&self) -> Result<&[u8], RawBuildInfoError> {
let build_info = self.build_info()?;
Ok(&self.data[..(build_info.bss_start - self.base_address()) as usize])
}
pub fn full_data(&self) -> &[u8] {
&self.data
}
pub fn base_address(&self) -> u32 {
self.offsets.base_address
}
pub fn end_address(&self) -> Result<u32, RawBuildInfoError> {
let build_info = self.build_info()?;
Ok(build_info.bss_end)
}
pub fn entry_function(&self) -> u32 {
self.offsets.entry_function
}
pub fn build_info_offset(&self) -> u32 {
self.offsets.build_info
}
pub fn autoload_callback(&self) -> u32 {
self.offsets.autoload_callback
}
pub fn overlay_signatures_offset(&self) -> u32 {
self.offsets.overlay_signatures
}
pub fn bss(&self) -> Result<Range<u32>, RawBuildInfoError> {
let build_info = self.build_info()?;
Ok(build_info.bss_start..build_info.bss_end)
}
pub fn offsets(&self) -> &Arm9Offsets {
&self.offsets
}
pub fn originally_compressed(&self) -> bool {
self.originally_compressed
}
pub fn originally_encrypted(&self) -> bool {
self.originally_encrypted
}
pub(crate) fn update_overlay_signatures(
&mut self,
arm9_overlay_table: &OverlayTable,
) -> Result<(), Arm9OverlaySignaturesError> {
let arm9_overlays = arm9_overlay_table.overlays();
let Some(signatures) = self.overlay_signatures_mut(arm9_overlays.len())? else {
return Ok(());
};
for overlay in arm9_overlays {
if let Some(signature) = overlay.signature() {
signatures[overlay.id() as usize] = signature;
}
}
if let Some(signature) = arm9_overlay_table.signature() {
let Some(table_signature) = self.overlay_table_signature_mut()? else {
return Ok(());
};
*table_signature = signature;
}
Ok(())
}
}
impl AsRef<[u8]> for Arm9<'_> {
fn as_ref(&self) -> &[u8] {
&self.data
}
}