path: root/srb_lib.py

#!/usr/bin/env python3
# Startdate: 2024-03-08-6 15:28
# File: srb_lib.py
# Purpose: library for srb.exe savegame hacking
# SPDX-License-Identifier: GPL-3.0-only
# History:
# Usage:
#    from srb.py
# Reference:
#    blog posts 2024-03
#    <https://github.com/8051Enthusiast/delsum>
#    crc width=32 poly=0x4c11db7 init=0x0 xorout=0x235b4b9c refin=false refout=false out_endian=little
#    Incorrect functions: https://gist.github.com/djsweet/3477115595efab31905be7000bb013bc FAILED
#    Fuctions that worked: https://gist.github.com/Lauszus/6c787a3bc26fea6e842dfb8296ebd630 Author: Lauszus
#    https://stackoverflow.com/questions/46109815/reversing-the-byte-order-of-a-string-containing-hexadecimal-characters
#    https://gamefaqs.gamespot.com/pc/930591-snoopy-vs-the-red-baron/faqs/46161
# Documentation:
#    winetricks vd=1024x768
#    winetricks vd=off
# Dependencies:
#    python3-crcmod

import sys, struct
srb_lib_version = "20240309a"

# Table of byte positions of values in the savegame file, minus the first four bytes which are the checksum. Due to zero-indexing of python lists, but for ease of usage, we will always put a zero as the value of index 0. That is, profile 1 will use index 1 of the list.
# money is 0x270 bytes after the "Z<dddddddd" start.
PROFILE_START_POSITION = [0, 0x10, 0x142C, 0x2848]
POS_MONEY = 0x270
POS_EQUIPPED_WEAPON = 0x284

WEAPONS = [
   "undefined", # 0x0
   "undefined", # secondary weapon is a single machine gun?!
   "undefined",
   "undefined",
   "undefined",
   "Potato Gun", # 0x5
   "Stinger",
   "Woodstock Missile",
   "Water Balloon Cannon",
   "Snow Blower",
   "Fire Boomerang", # 0xA
   "Lightning Rod",
   "Bottle Rockets",
   "Roman Candles",
   "10 Gauge Pumpkin",
   "undefined", # 0xF
]

# DEFINE FUNCTIONS
def ferror(*args, **kwargs):
   print(*args, file=sys.stderr, **kwargs)

def debuglev(_numbertocheck, debuglevel):
   # if _numbertocheck <= debuglevel then return truthy
   _debuglev = False
   try:
      if int(_numbertocheck) <= int(debuglevel):
         _debuglev = True
   except Exception as e:
      pass
   return _debuglev

def read_file(filename):
   """ Return the stored checksum and contents of the file """
   with open(filename, 'rb') as f:
      # Read the entire file
      file_contents = f.read()
      checksum = file_contents[0:4]
      file_contents_without_checksum = file_contents[4:]
   return checksum, file_contents_without_checksum

def write_file(filename, checksum, data):
   """ Write the checksum and data to the file. We are assuming the checksum is already correct. """
   print(f"debug: checksum(type={type(checksum).__name__})={checksum}")
   if type(checksum) == int:
      checksum = checksum
   elif type(checksum) == str:
      checksum = int(checksum)
   elif type(checksum) == bytes:
      checksum = int.from_bytes(checksum)
   with open(filename, 'wb') as f:
      f.write(struct.pack('>L',checksum))
      #f.write(struct.pack('>I',int(checksum)))
      f.write(data)

######## start copy-paste 2
#  https://gist.github.com/Lauszus/6c787a3bc26fea6e842dfb8296ebd630

def reflect_data(x, width):
   # See: https://stackoverflow.com/a/20918545
   if width == 8:
      x = ((x & 0x55) << 1) | ((x & 0xAA) >> 1)
      x = ((x & 0x33) << 2) | ((x & 0xCC) >> 2)
      x = ((x & 0x0F) << 4) | ((x & 0xF0) >> 4)
   elif width == 16:
      x = ((x & 0x5555) << 1) | ((x & 0xAAAA) >> 1)
      x = ((x & 0x3333) << 2) | ((x & 0xCCCC) >> 2)
      x = ((x & 0x0F0F) << 4) | ((x & 0xF0F0) >> 4)
      x = ((x & 0x00FF) << 8) | ((x & 0xFF00) >> 8)
   elif width == 32:
      x = ((x & 0x55555555) << 1) | ((x & 0xAAAAAAAA) >> 1)
      x = ((x & 0x33333333) << 2) | ((x & 0xCCCCCCCC) >> 2)
      x = ((x & 0x0F0F0F0F) << 4) | ((x & 0xF0F0F0F0) >> 4)
      x = ((x & 0x00FF00FF) << 8) | ((x & 0xFF00FF00) >> 8)
      x = ((x & 0x0000FFFF) << 16) | ((x & 0xFFFF0000) >> 16)
   else:
      raise ValueError('Unsupported width')
   return x

def crc_poly(data, n, poly, crc=0, ref_in=False, ref_out=False, xor_out=0):
   g = 1 << n | poly  # Generator polynomial
   # Loop over the data
   for d in data:
      # Reverse the input byte if the flag is true
      if ref_in:
         d = reflect_data(d, 8)
      # XOR the top byte in the CRC with the input byte
      crc ^= d << (n - 8)
      # Loop over all the bits in the byte
      for _ in range(8):
         # Start by shifting the CRC, so we can check for the top bit
         crc <<= 1
         # XOR the CRC if the top bit is 1
         if crc & (1 << n):
            crc ^= g
   # Reverse the output if the flag is true
   if ref_out:
      crc = reflect_data(crc, n)
   # Return the CRC value
   return crc ^ xor_out
##### end copy-paste 2

def _get_data_from_data_object(data_object):
   """ Helper function to either open file, or pass bytes through. """
   if type(data_object) == str:
      _, data = read_file(data_object)
   elif type(data_object) == bytes:
      data = data_object
   return data

def _convert_dec_to_reverse_bytes(num_dec):
   """ Helper function to turn a given decimal value into the little-endian values for insertion into the data. """
   # Here is how to do this manually.
   #num_str = str(hex(num_dec))[2:]
   #num_str = ("0" if len(num_str) % 2 else "") + num_str
   #num_bytes = bytearray.fromhex(num_str)
   #num_bytes.reverse()
   #return bytes(num_bytes)
   # I learned how to do this the pythonic way.
   return struct.pack('<i',num_dec)

def get_money(data_object,profile_id):
   """ Print in decimal the current money value of the profile_id. """
   data = _get_data_from_data_object(data_object)
   money_dec = data[PROFILE_START_POSITION[profile_id]+POS_MONEY] + (data[PROFILE_START_POSITION[profile_id]+POS_MONEY+1] * 0x100)
   #money_dec = int(money_hex, base=16)
   return money_dec

def set_money(data_object,profile_id, money_dec):
   """ Using data, set the money given in decimal for the given profile_id. """
   data = _get_data_from_data_object(data_object)
   #money_bytes = _convert_dec_to_reverse_bytes(money_dec)
   #pos_money_low = PROFILE_START_POSITION[profile_id]+POS_MONEY
   data_bytearray = bytearray(data)
   struct.pack_into('<i',data_bytearray,PROFILE_START_POSITION[profile_id]+POS_MONEY,money_dec)
   data = bytes(data_bytearray)
   print(f"after setting money to {money_dec}, we checked and got {get_money(data,profile_id)}")
   return data

def get_weapon(data_object,profile_id):
   data = _get_data_from_data_object(data_object)
   weapon = data[PROFILE_START_POSITION[profile_id]+POS_EQUIPPED_WEAPON]
   return WEAPONS[weapon]

def set_weapon(data_object,profile_id,weapon):
   data = _get_data_from_data_object(data_object)
   data_bytearray = bytearray(data)
   weapon, weapon_name = get_weapon_info(weapon)
   if weapon == -1:
      ferror(f"Warning: cannot set weapon to {weapon}, because {weapon_name}")
      return -1
   # armed with weapon as the index number, let's change the savegame data
   struct.pack_into('<i',data_bytearray,PROFILE_START_POSITION[profile_id]+POS_EQUIPPED_WEAPON,weapon)
   data = bytes(data_bytearray)
   print(f"after setting weapon to {WEAPONS[weapon]}, we checked and got {get_weapon(data,profile_id)}")
   return data

def get_weapon_info(weapon):
   """ Returns index, and name of weapon, or -1 and "invalid" """
   if type(weapon) == str:
      try:
         weapon_id = [index for index in range(len(WEAPONS)) if WEAPONS[index].lower() == weapon.lower()][0]
         return weapon_id, WEAPONS[weapon_id]
      except:
         return -1, f"cannot find weapon {weapon}" # must be an incorrect name
   elif type(weapon) == int:
      if weapon not in range(0,16):
         return -1, f"invalid index {weapon}; use 0-15" # must be <0 or >15
      else:
         return weapon, WEAPONS[weapon]
   else:
      return -1, f"invalid way to reference weapon: [{type(weapon)}]. Use index or name."

def calculate_checksum(data):
   """ Return the 4-byte checksum used by the game for the provided data. """
   # aka # CRC-32/BZIP2
   # The polynomial appears to be a standard one like used for Ethernet, but I am uncertain of the xor is standard or not.
   # I learned all this only with <https://github.com/8051Enthusiast/delsum>
   # crc width=32 poly=0x4c11db7 init=0x0 xorout=0x235b4b9c refin=false refout=false out_endian=little
   csum = hex(crc_poly(data, 32, 0x4C11DB7, crc=0x0, ref_in=False, ref_out=False, xor_out=0x235b4b9c))
   # comes in as ceb434d4, but needs to be d434b4ce, and bytearray is convenient way to swap it like that
   # trim off the '0x' from the string
   #print(f"Got csum={csum}")
   csum = bytearray.fromhex(str(csum)[2:])
   csum.reverse()
   # and back to bytes because that is how we will want to use it.
   return bytes(csum)

def correct_file(filename, debuglevel = 0):
   """ Given a savegame file, calculate the correct checksum and store that sum instead of whatever is there. """
   cs, data = read_file(filename)
   csum = calculate_checksum(data)
   ran = False
   if csum != cs:
      if debuglev(5,debuglevel):
         ferror(f"Stored checksum is {cs} and will update it to {csum}")
      write_file(filename,csum,data)
      ran = True
   else:
      if debuglev(2,debuglevel):
         ferror(f"Stored checksum is still correct, {csum}. Skipping {filename}")
   return ran