rusty_uma_extractor 0.1.0

use std::{
    cmp,
    error::{self, Error},
    fmt, 
    fs::OpenOptions, 
    io::{ 
        self, BufRead, BufReader, Cursor, IoSliceMut, Read, Seek, SeekFrom 
    },
    num,
    path::Path,
    sync::LazyLock
};

use log::{self, kv};

use rmp::{ self, decode::DecodeStringError };
use frida::Frida;

static FRIDA: LazyLock<Frida> = LazyLock::new(|| unsafe { Frida::obtain() });

type Result<T> = std::result::Result<T,VetExtractorError>;

#[derive(Debug)]
pub enum VetExtractorError {
    /// Raised when `frida::DeviceManager` could not get a lock on device. See `frida` crate for possible reasons
    GetDeviceFail,
    /// Raised when the process for _Uma Musume: Pretty Derby_ could not be found in running processes list for target device
    ProcessNotFound,
    /// Wrapper for `std::io::Error`
    Io(io::Error),
    /// Wrapper for `std::num::ParseIntError`
    ParseInt(num::ParseIntError),
    /// Wrapper for `std::num::TryFromIntError`
    TryFromInt(num::TryFromIntError)
    
}

impl fmt::Display for VetExtractorError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            VetExtractorError::GetDeviceFail =>
                write!(f,"Failed to get device"),
            VetExtractorError::ProcessNotFound =>
                write!(f,"Could not find active process"),
            VetExtractorError::Io(err) =>
                write!(f,"I/O Error: {:#?}",err),
            VetExtractorError::ParseInt(err) =>
                write!(f,"Encountered ParseIntError: {:#?}",err),
            VetExtractorError::TryFromInt(err) =>
                write!(f,"Encountered TryFromIntError: {:#?}",err)
        }
    }
}

impl error::Error for VetExtractorError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            VetExtractorError::GetDeviceFail => None,
            VetExtractorError::ProcessNotFound => None,
            VetExtractorError::Io(io_err) => Some(io_err),
            VetExtractorError::ParseInt(int_err) => Some(int_err),
            VetExtractorError::TryFromInt(int_err) => Some(int_err)
        }
    }
}

/// Find running instance process and return PID
pub fn find_process() -> Result<u32> {
    let device_manager = frida::DeviceManager::obtain(&FRIDA);
    let Ok(device) = device_manager.get_local_device() else {
        log::error!(target: "find_process", "frida::DeviceManager could not obtain lock on target device.");
        return Err(VetExtractorError::GetDeviceFail)
    };

    let processes = device.enumerate_processes();

    let Some(process) = processes.iter().find(|p| p.get_name().contains("UmamusumePrettyDerby.exe")) else {
        log::error!(target: "find_process", "Could not locate active process for \"UmamusumePrettyDerby.exe\"");
        return Err(VetExtractorError::ProcessNotFound)
    };
    let proc_id = process.get_pid();
    let proc_name = process.get_name();
    log::debug!(target: "find_process",proc_id, proc_name; "Found process");
    Ok(proc_id)
}

/// Object structure for a mapped memory region of the running instance

#[derive(Debug, Copy, Clone)]
pub struct MemRegion {
    /// Start address of memory block
    pub start_addr: u64,
    /// End address of memory block
    pub end_addr: u64,
    /// Size of memory block
    pub size: usize
}

impl fmt::Display for MemRegion {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let size = u32::try_from(self.size).unwrap_or_default();
        let f_size = f64::from(size);

        let fmt_size = match f_size {
            s if s <= 1024f64 => (s, "B"),
            s if s <= 1024f64 * 1024f64 => (s / 1024f64,"KB"),
            s if s < 1024f64 * 1024f64 * 1024f64 => (s / 1024f64 / 1024f64, "MB"),
            _ => (f_size, "B")
        };

        write!(
            f, "MemRegion\n\t{:x} - {:x}\n\tSize: {:.5} {}",
            self.start_addr, self.end_addr, fmt_size.0, fmt_size.1
        )

    }
}

/// Read memory maps to find RW regions the instance uses to store trained character data
/// 
/// _**Note**: This function requires read-only permission access for `/proc/{PID}/maps` in order to map the active memory regions
/// currently in use by the process._
/// 
/// # Params
/// | name | type | description |
/// | -- | -- | -- |
/// | `pid` | `u32` | PID of running process found for instance on target device. |
pub fn get_memmap(pid: u32) -> Result<Vec<MemRegion>> {
    let path_str = format!("/proc/{}/maps",pid);

    let maps_path = Path::new(path_str.as_str());

    log::debug!(target: "get_memmap","opening file for: {:#?}",&maps_path);

    // Attempt to open maps file for found PID 
    let maps_file = match OpenOptions::new().read(true).open(maps_path) {
        Ok(f) => f,
        Err(e) => {
            let e_msg = format!("Failed to open file: {:#?}\nError: {:#?}", maps_path, e.to_string());
            log::error!(target: "get_memmap", "{}",e_msg);
            return Err(VetExtractorError::Io(e))
        }
    };

    // BufReader to read file and Vec to hold mapped mem regions we're interested in.
    let mut reader = BufReader::new(maps_file);
    let mut regions:Vec<MemRegion> = Vec::new();

    log::debug!(target: "get_memmap" ,"reading mem map addresses...");
    let mut buf = String::new();
    // continue to read lines in file until read_line returns Ok(0) meaning EOF
    while let Ok(read_bytes) = reader.read_line(&mut buf) {
        if read_bytes > 0 {
            let mut parts = buf.split_ascii_whitespace();
            // Parse address range and size
            let (start_addr, end_addr, size) = match parts.next() {
                Some(range) => {
                    let addr: Vec<&str> = range.split("-").collect();

                    let start_addr: u64 = match u64::from_str_radix(addr[0],16) {
                        Ok(n) => n,
                        Err(e) => return Err(VetExtractorError::ParseInt(e))
                    };
                    let end_addr: u64 = match u64::from_str_radix(addr[1],16){
                        Ok(n) => n,
                        Err(e) => return Err(VetExtractorError::ParseInt(e))
                    };
                    let size: usize = match (&end_addr - &start_addr).try_into() {
                        Ok(n) => n,
                        Err(e) => return Err(VetExtractorError::TryFromInt(e))
                    };

                    (start_addr, end_addr, size)
                },
                None => (0u64,0u64,0)
            };

            // Only interested in rw regions
            let is_rw: bool = match parts.next() {
                Some(perms) => {
                    perms.contains("rw")
                },
                None => false
            };

            // If region is rw and size between <100 KB and 500 MB, push to Vec
            if is_rw && (size >= 100_000 && size <= 512 * 1024 * 1024) {
                regions.push(MemRegion { start_addr, end_addr, size });
            }

            buf.clear();
        } else {
            break
        }
    }

    // Sort found regions by descending size, data likely in larger allocations
    regions.sort_by(
        |a, b| b.size.cmp(&a.size)
    );

    log::debug!(target: "get_memmap","Complete. Found: {:?} regions",regions.len());
    Ok(regions)
}

/// Object structure for a region of memory that is a possible candidate to contain trained character data

#[derive(Debug, Copy, Clone)]
pub struct TrainedCharacterArray {
    /// MemRegion object
    pub mem_region: MemRegion,
    /// Offset value for start of trained character array in region
    pub arr_offset: u64,
    /// Number of elements found in trained character array in region
    pub n_elements: u32
}

/// Searches memory regions found in-use by active instance in attempt to find trained character array
/// 
/// _**Note**: This function requires read-only permission access for `/proc/{PID}/mem` in order to map the active memory regions
/// currently in use by the process._
/// 
/// # Params
/// | name | type | description |
/// | -- | -- | -- |
/// | `pid` | `u32` | PID of running process found for instance on target device. |
/// | `scan_regions` | `Vec<MemRegion>` | Vector array of `MemRegion` objects in which to search for trained character array data. |

pub fn chara_array_data_search(pid: u32, scan_regions: Vec<MemRegion>) -> Result<Option<TrainedCharacterArray>> {

    const TRAINED_CHARA_ARRAY_FIXSTR: &str = "trained_chara_array";
    const U8_FIXSTR_MARKER: rmp::Marker = rmp::Marker::FixStr(19);

    let path_str = format!("/proc/{}/mem",pid);

    let mem_path = Path::new(path_str.as_str());
    log::debug!(target: "chara_array_data_search","opening file for: {:?}",&mem_path);
    // Attempt to open mem file for found PID 
    let mem_file = match OpenOptions::new().read(true).open(mem_path) {
        Ok(f) => f,
        Err(e) => {
            let e_msg = format!("Failed to open file: {:#?}\nError: {:#?}", mem_path, e.to_string());
            log::error!(target: "get_memmap", "{}",e_msg);
            return Err(VetExtractorError::Io(e))
        }
    };

    // BufReader to read file
    let mut reader = BufReader::new(mem_file);

    log::info!(target: "chara_array_data_search","Beginning search in {:?} regions...",scan_regions.len());

    for region in scan_regions.iter() {
        log::debug!(target: "chara_array_data_search", "reading region: {}",region);
        let s_addr = region.start_addr;
        log::debug!(target: "chara_array_data_search" ,"seeking to region start: {:x?}",s_addr);
        match reader.seek(SeekFrom::Start(region.start_addr)) {
            Ok(_) => (),
            Err(e) => {
                let e_msg = format!(
                    "Buffer failed SeekFrom operation to region start addr: {:x?}\nError: {}",
                    region.start_addr,
                    e.to_string()
                );
                log::error!(target: "chara_array_data_search", "{}", e_msg);
                return Err(VetExtractorError::Io(e))
            }
        };
        
        let mut buf = vec![0u8;region.size];
        match reader.read_exact(&mut buf) {
            Ok(_) => (),
            Err(e) => return Err(VetExtractorError::Io(e))
        };

        let mut cur = Cursor::new(&mut buf);
        let mut bytes_read: usize = 0;

        while let Ok(skipped) = cur.skip_until(U8_FIXSTR_MARKER.to_u8()) {
            bytes_read += skipped;

            // Ensure we haven't reached the end before trying to read
            let buf_len = match u64::try_from(cur.get_ref().len()) {
                Ok(len) => len,
                Err(e) => {
                    log::error!(target: "chara_array_data_search", "u64::TryFrom<usize> operation failed: {}",e.to_string());
                    return Err(VetExtractorError::TryFromInt(e))
                }
            };
            if cur.position() == buf_len {
                let cur_end_pos = cur.position();
                log::debug!(
                    target: "chara_array_data_search",
                    cur_end_pos,
                    buf_len,
                    bytes_read;
                    "Search reached end of region. Yielding search to next region..."
                );
                break
            }

            // Move cursor back 1 pos in order to read the marker in
            match cur.seek(SeekFrom::Current(-1)) {
                Ok(_) => (),
                Err(e) => {
                    log::error!(target: "chara_array_data_search","SeekFrom::Current(-1) operation failed: {}",e.to_string());
                    return Err(VetExtractorError::Io(e))
                }
            };
            bytes_read -= 1;

            let cur_position: usize = match usize::try_from(cur.position()) {
                Ok(n) => n,
                Err(e) => {
                    log::error!(target: "chara_array_data_search", "usize::TryFrom<u64> operation failed: {}",e.to_string());
                    return Err(VetExtractorError::TryFromInt(e))
                }
            };
            let end_offset: usize = cur_position + 23;

            let mut str_buff = [0u8;20];
            let mut arr16_buff = [0u8;3];

            // Create array of slices to read into
            let str_io_slice  = IoSliceMut::new(&mut str_buff);    // slice for "trained_chara_array" fixstr
            let arr16_io_slice = IoSliceMut::new(&mut arr16_buff);  // slice for arr16 len

            let read_size = match cur.read_vectored(&mut [str_io_slice,arr16_io_slice]) {
                Ok(bytes) => bytes,
                Err(e) => {
                    log::error!(target: "chara_array_data_search", "Cursor failed to read into vectored buffers: {}",e.to_string());
                    return Err(VetExtractorError::Io(e))
                }
            };
            bytes_read += read_size;

            // Expected len of data read is 23 bytes consisting of:
            // FixStr(19) Marker(0xb3) = 1 byte; "trained_chara_array"b = 19 bytes
            // Arr16 Marker(0xdc) = 1 byte; YYYYYYYY YYYYYYYY = 16-bit (2 bytes) big-endian uint represents size of array elements
            if read_size != 23 {
                continue
            }

            // Expect only sub-set of Basic Latin UTF-8 characters in string search
            // Guard against noise
            let utf_range: std::ops::RangeInclusive<u8> = 0x20..=0x7f;
            if str_buff[1..].iter().any(| b| !utf_range.contains(b)) {
                continue
            }

            // decode the bytes into string using a buffer of FixStr len: 19
            let mut decoded_str_buf = [0u8;19];

            let decode_str = match rmp::decode::read_str(&mut &str_buff[..], &mut decoded_str_buf[..]) {
                    Ok(str) => str,
                    Err(e) => {
                        log::error!(target: "chara_array_data_search", "MessagePack failed to decode string bytes: {}",e.to_string());
                        continue
                    }
                };

            if decode_str == TRAINED_CHARA_ARRAY_FIXSTR {
                // If string is found, check arr16 slice
                let marker_byte = arr16_buff[0];

                if let rmp::Marker::Array16 = rmp::Marker::from_u8(marker_byte) {
                    log::debug!(target: "chara_array_data_search","Found Array16 marker! Getting number of elements...");

                    // Get number of elements in arr16 data block to expect
                    let n_elements = match rmp::decode::read_array_len(&mut &arr16_buff[..]) {
                        Ok(n) => {
                            log::debug!(target: "chara_array_data_search",n ;"Number of array elements");
                            n
                        },
                        Err(e) => {
                            log::error!(target: "chara_array_data_search", "Failed to decode array len: {}",e.to_string());
                            0
                        }
                    };
                    // Set array offset value to the start of the arr16 marker
                    let arr_offset = match u64::try_from(end_offset - 3) {
                        Ok(n) => n,
                        Err(e) => {
                            log::error!(target: "chara_array_data_search", "u64::TryFrom<usize> operation failed: {}", e.to_string());
                            return Err(VetExtractorError::TryFromInt(e))
                        }
                    };
                    let start_addr = region.start_addr;
                    log::info!(
                        target: "chara_array_data_search",
                        start_addr,
                        arr_offset,
                        n_elements;
                        "Found trained character data in region"
                    );

                    return Ok(Some( TrainedCharacterArray { mem_region: *region, arr_offset, n_elements } ))

                }

            }

        }

    }

    log::info!(target: "chara_array_data_search","No card_id entries found in mem regions");
    Ok(None)
}

pub type TrainedCharacterData = Vec<u8>;

/// Attempts to extract the Arr16 byte-array containing the trained character data
/// 
/// _**Note**: This function requires read-only permission access for `/proc/{PID}/mem` in order to map the active memory regions
/// currently in use by the process._
/// 
/// # Params
/// | name | type | description |
/// | -- | -- | -- |
/// | `pid` | `u32` | PID of running process found for instance on target device. |
/// | `character_arr` | `TrainedCharacterArray` | The `TrainedCharacterArray` object found as best candidate to contain the trained character data |

pub fn get_chara_array(pid: u32 ,character_arr: TrainedCharacterArray) -> Result<Option<TrainedCharacterData>> {
    let TrainedCharacterArray { mem_region, arr_offset, n_elements } = character_arr;

    let path_str = format!("/proc/{}/mem",pid);

    let mem_path = Path::new(path_str.as_str());
    log::info!(target: "get_chara_array", "Opening file for: {:?}",&mem_path);
    // Attempt to open maps file for found PID 
    let mem_file = match OpenOptions::new().read(true).open(mem_path) {
        Ok(f) => f,
        Err(e) => {
            let e_msg = format!("Failed to open file: {:#?}\nError: {:#?}", mem_path, e.to_string());
            log::error!(target: "get_chara_array", "{}",e_msg);
            return Err(VetExtractorError::Io(e))
        }
    };

    // BufReader to read file
    let mut reader = BufReader::new(mem_file);

    log::debug!(target: "get_chara_array","Seeking to start of memory region: {:x?}",mem_region.start_addr);
    match reader.seek(SeekFrom::Start(mem_region.start_addr)) {
        Ok(_) => (),
        Err(e) => {
            let e_msg = format!(
                    "Buffer failed SeekFrom operation to region start addr: {:x?}\nError: {}",
                    mem_region.start_addr,
                    e.to_string()
                );
                log::error!(target: "chara_array_data_search", "{}", e_msg);
            return Err(VetExtractorError::Io(e))
        }
    };

    log::debug!(target: "chara_array_data_search","Reading region into buffer: {}",mem_region);

    let mut buf = vec![0u8;mem_region.size];
    match reader.read_exact(&mut buf) {
        Ok(_) => (),
        Err(e) => {
            log::error!(target: "get_chara_array", "Reader failed to read into buffer: {}",e.to_string());
            return Err(VetExtractorError::Io(e))
        }
    };
    
    let mut cur = Cursor::new(&mut buf);
    log::debug!(target: "get_chara_array","Cursor currently at {:?} position, moving to array start at: {:?}",cur.position(),arr_offset);

    cur.set_position(arr_offset);
    log::debug!(target: "get_chara_array","Cursor now at array start: {:?}",cur.position());

    let mut arr16_buff = [0u8;3];
    let arr16_io_slice = IoSliceMut::new(&mut arr16_buff);

    match cur.read_vectored(&mut [arr16_io_slice]) {
        Ok(_) => (),
        Err(e) => {
            log::error!(target: "get_chara_array","Reader failed to read into vectored buffers: {}",e.to_string());
            return Err(VetExtractorError::Io(e))
        }
    };

    let marker_byte = arr16_buff[0];

    let is_arr: bool = match rmp::Marker::from_u8(marker_byte) {
        rmp::Marker::Array16 => match rmp::decode::read_array_len(&mut &arr16_buff[..]) 
            {
                Ok(n) => {
                    if n == n_elements {
                        log::info!(target: "get_chara_array", n; "Number of trained character records in array");
                        true
                    } else {
                        log::warn!(
                            target: "get_chara_array",
                            found = n, expected = n_elements;
                            "Number of records found did not match expected value"
                        );
                        false
                    }
                },
                Err(e) => {
                    log::error!(target: "get_chara_array","Failed to decode array len: {:?}",e);
                    false
                }
            }
        ,
        m@ _ => {
            log::warn!(target: "get_chara_array","Found unexpected marker: {:?}",m);
            false
        }
    };

    if is_arr {
        // Move cursor back to start of array
        let _ = cur.seek(SeekFrom::Current(-3));
        let mut found_data: Option<TrainedCharacterData> = None;
        /* 
        *   This section iterates through possible block sizes and then
        *   performs a "sanity check" by counting the number of instances of "card_id" keys.
        *   
        *   The expectation is that for each trained character record, there should be 7 instances
        *   of "card_id": trained char + 2 parents + 4 grandparents.
        *
        *   For performance, we only want to check a max of 3MB for card_id instances.
        */
        let sizes: [u64; 3] = [15 * 1024 * 1024, 20 * 1024 * 1024, 25 * 1024 * 1024];

        for sz in 0..3 {
            let mem_region_size = match u64::try_from(mem_region.size) {
                Ok(n) => n,
                Err(e) => return Err(VetExtractorError::TryFromInt(e))
            };
            let search_size: u64 = mem_region_size - arr_offset;
            let max_read_size = cmp::min(sizes[sz],search_size);

            if max_read_size < 1024u64 * 1024u64 {
                continue
            }
            
            let buf_size: usize = match usize::try_from(max_read_size) {
                Ok(n) => n,
                Err(e) => return Err(VetExtractorError::TryFromInt(e))
            };
            let mut potential_data = vec![0u8;buf_size];

            match cur.read_exact(&mut potential_data) {
                Ok(_) => (),
                Err(e) => {
                    log::error!(
                        target: "get_chara_array",
                        "Reader failed reading potential data into buffer with size: {:?}\n{}",
                        buf_size, e.to_string()
                    );
                    cur.set_position(arr_offset);
                    continue
                }
            }

            let mut data_cur = Cursor::new(&mut potential_data);
            let card_marker: rmp::Marker = rmp::Marker::FixStr(7);
            let search_str = "card_id";
            let check_size = cmp::min(buf_size - 8,3 * 1024 * 1024);
            let mut read_bytes: usize = 0;
            let mut card_count: usize = 0;
            let san_check: usize = match usize::try_from(n_elements) {
                Ok(n) => n,
                Err(e) => {
                    log::error!(target: "get_chara_array", "usize::TryFrom<u64> operation failed: {}",e.to_string());
                    return Err(VetExtractorError::TryFromInt(e))
                }
            };

            while let Ok(skipped) = data_cur.skip_until(card_marker.to_u8()) {
                // Cursor is moved back 1 position here so decode operation can
                //  read in marker byte
                match data_cur.seek(SeekFrom::Current(-1)) {
                    Ok(_) => (),
                    Err(e) => {
                        log::error!(target: "get_chara_array", "SeekFrom::Current(-1) operation failed: {}",e.to_string());
                        return Err(VetExtractorError::Io(e))
                    }
                };
                read_bytes += skipped - 1;

                let mut str_buf = [0u8;7];
                let mut in_buf = [0u8;8];
                match data_cur.read_exact(&mut in_buf) {
                    Ok(_) => (),
                    Err(e) => {
                        log::error!(target: "get_chara_array", "Reader failed to read into buffer: {}",e.to_string());
                        return Err(VetExtractorError::Io(e))
                    }
                };

                let res_str = match rmp::decode::read_str(&mut &in_buf[..],&mut str_buf[..]) {
                    Ok(str) => str,
                    Err(_) => continue
                };
                read_bytes += 8;
                if res_str == search_str {
                    card_count += 1;
                    // Copied from original Python logic
                    if card_count >= 200 {
                        break
                    }
                }

                // Break when end of the check size reached
                if read_bytes >= check_size {
                    break
                }
            }
            /*
            *   Original logic dictated that if card_count >= 150, the correct region was most likely found,
            *   however - if a user had less than 22 trained characters, this would not trigger.
            *
            *   Logic was updated so as to either break if more than 200 cards were found OR if card_count / n_elements == 7
            *   which covers the case of users having less than 22 trained characters.
             */

            if card_count >= 200 || card_count / san_check == 7 {
                log::info!(target: "get_chara_array","Found array with {:?}+ cards",card_count);
                found_data = Some(potential_data);
                break
            }

        }
        return Ok(found_data)
    }
    Ok(None)
}