#!/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:

import sys, struct
srb_lib_version = "20240311a"

# 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
POS_PROFILE_IN_USE = 0x290
POS_NAME = 0x294

NAME_CHARS = " ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890!?"

POS_BYTES_WEAPONS_PURCHASED = 0x028C
# Woodstock Missile and Bottle Rockets are always available to use.
# for equipped-weapon
# id = value of equipped byte
# e = equippable
# p = purchase flag, where in the two bytes (when looking at it big-endian) the entry is.
WEAPONS = [
   {"id":0x0,"e":False,"p":0x0,   "name":"none"}, # custom for this library
   {"id":0x1,"e":False,"p":0x0,   "name":"machine-gun-single"}, # not part of the game, but works if you set manually
   {"id":0x2,"e":False,"p":0x0,   "name":"machine-gun-double"}, # not part of the game, but works if you set manually
   {"id":0x3,"e":False,"p":0x0,   "name":"machine-gun-triple"}, # not part of the game, but works if you set manually
   {"id":0x4,"e":False,"p":0x0,   "name":"machine-gun-four"},
   {"id":0x5,"e":True, "p":0x0020,"name":"Potato Gun"}, # always available to buy
   {"id":0x6,"e":True, "p":0x0040,"name":"Stinger"}, # always available to buy
   {"id":0x7,"e":True, "p":0x0,   "name":"Woodstock Missile"}, # always available to player
   {"id":0x8,"e":True, "p":0x0100,"name":"Water Balloon Cannon"},
   {"id":0x9,"e":True, "p":0x0200,"name":"Snow Blower"},
   {"id":0xA,"e":True, "p":0x0400,"name":"Fire Boomerang"},
   {"id":0xB,"e":True, "p":0x0800,"name":"Lightning Rod"},
   {"id":0xC,"e":True, "p":0x0,   "name":"Bottle Rockets"}, # always available to player
   {"id":0xD,"e":True, "p":0x2000,"name":"Roman Candles"},
   {"id":0xE,"e":True, "p":0x4000,"name":"10 Gauge Pumpkin"},
   {"id":0xF,"e":False,"p":0x6F60,"name":"all"}, # custom for this library
]

# WORKHERE: position is for where the balloon info is stored for this levelset, if it is stored here and not per-level.
LEVELSETS = [
   {"id":0,"name":"Aerodrome Island","pos":0x0},
   {"id":1,"name":"Woods of Montsec","pos":0x0},
   {"id":2,"name":"Front Lines of Verdon","pos":0x0},
   {"id":3,"name":"Mines of the Matterhorn","pos":0x0},
   {"id":4,"name":"Verdon Gorge","pos":0x0},
   {"id":5,"name":"Flying Fortress","pos":0x0},
]

# relative to player profile
POS_LEVEL_START = 0x2D8
# each level position is POS_LEVEL_START + (4*level["pos_r"])
# the relative position is because each level set gets 7 (or 6 plus 1 blank) slots, but not all levelsets have 6 levels.
LEVELS = [
   {"id":0, "pos_r":0, "setid":0,"name":"Defend Island"},
   {"id":1, "pos_r":1, "setid":0,"name":"Recover Plans"},
   {"id":2, "pos_r":2, "setid":0,"name":"Protect the Trucks"},
   {"id":3, "pos_r":3, "setid":0,"name":"Attack of the U-Boats"},
   {"id":4, "pos_r":4, "setid":0,"name":"Cripple Outpost Island"},
   {"id":5, "pos_r":5, "setid":0,"name":"Sink the Battleship"},
   {"id":6, "pos_r":7, "setid":1,"name":"Rerun's Challenge"},
   {"id":7, "pos_r":8, "setid":1,"name":"Eliminate Tree Village"},
   {"id":8, "pos_r":9, "setid":1,"name":"Tree Chopper"},
   {"id":9, "pos_r":14,"setid":2,"name":"Trench Warfare"},
   {"id":10,"pos_r":15,"setid":2,"name":"Recover Allied Base"},
   {"id":11,"pos_r":16,"setid":2,"name":"Giant Tank"},
   {"id":12,"pos_r":21,"setid":3,"name":"Surprise Attack"},
   {"id":13,"pos_r":22,"setid":3,"name":"Derail the Train"},
   {"id":14,"pos_r":23,"setid":3,"name":"Enter the Mines"},
   {"id":15,"pos_r":24,"setid":3,"name":"Explore the Mines"},
   {"id":16,"pos_r":25,"setid":3,"name":"Destroy Driller Boss"},
   {"id":17,"pos_r":28,"setid":4,"name":"Navigate the Canyon"},
   {"id":18,"pos_r":29,"setid":4,"name":"Destroy Circus City"},
   {"id":19,"pos_r":30,"setid":4,"name":"Circus Aircraft Carrier"},
   {"id":20,"pos_r":35,"setid":5,"name":"Rescue Allies"},
   {"id":21,"pos_r":36,"setid":5,"name":"Battle the Red Baron"}
]

# hex values
# WOKRHERE: profile 2 level statuses are off somehow?! they are 4 bytes closer to front than the calculation shows.
LEVEL_STATUSES = {
   "59": "all-balloons,all-letters,sergeant",
   "5B": "all-balloons,all-letters,colonel",
   "5D": "all-balloons,all-letters,colonel",
   "60": "all-balloons,all-letters,colonel",
   "62": "all-balloons,all-letters,colonel",
   "63": "all-balloons,all-letters,colonel",
   "64": "all-balloons,all-letters,general"
}

POS_TUTORIAL_COMPLETED = 0x2D0 # bool

# 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. """
   data = None
   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)
   data = srb_pack('<i',data,PROFILE_START_POSITION[profile_id]+POS_MONEY,money_dec)
   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):
   """ For given profile, print currently equipped weapon. """
   data = _get_data_from_data_object(data_object)
   weapon_id = data[PROFILE_START_POSITION[profile_id]+POS_EQUIPPED_WEAPON]
   weapon_id = int(weapon_id)
   try:
      weapon, message = get_weapon_info(weapon_id)
   except:
      return f"unable to determine weapon"
   if message != "":
      return f"failed to get weapon due to {message}"
   print(f"debug: got weapon {weapon}")
   return weapon["name"]

def set_weapon(data_object,profile_id,weapon):
   data = _get_data_from_data_object(data_object)
   weapon, message = get_weapon_info(weapon)
   if (not bool(weapon)) or message != "":
      ferror(f"Warning: cannot set weapon to {weapon}, because {message}")
      return -1
   # armed with weapon as the index number, let's change the savegame data
   data = srb_pack('<i',data,PROFILE_START_POSITION[profile_id]+POS_EQUIPPED_WEAPON,weapon["id"])
   print(f"after setting weapon to {weapon['name']}, we checked and got {get_weapon(data,profile_id)}")
   return data

def get_weapon_info(weapon):
   """
   Returns tuple of (weapon object, message) searching by name or id. Message is empty if the weapon object is returned.
   """
   try:
      weapon = int(weapon)
   except:
      pass
   if type(weapon) == str:
      try:
         weapon_obj = [i for i in WEAPONS if i["name"].lower() == weapon.lower()][0]
         #weapon_id = [index for index in range(len(WEAPONS)) if WEAPONS[index].lower() == weapon.lower()][0]
         return weapon_obj, ""
      except:
         return {}, f"cannot find weapon {weapon}" # must be an incorrect name
   elif type(weapon) == int:
      if weapon in range(0,16):
         weapon_obj = [i for i in WEAPONS if i["id"] == weapon][0]
         return weapon_obj, ""
         #return weapon, WEAPONS[weapon]
      else:
         return {}, f"invalid index {weapon}; use 0-15" # must be <0 or >15
   return {}, f"invalid way to reference weapon: [{type(weapon)}]. Use index or name."

def get_purchased_weapons(data_object,profile_id,silent=False):
   """
   For the given profile, return which weapons are already purchased, as a comma-separated string and a bitmask.
   """
   data = _get_data_from_data_object(data_object)
   weapons_purchased = struct.unpack_from('<1i',data,PROFILE_START_POSITION[profile_id]+POS_BYTES_WEAPONS_PURCHASED)[0]
   #print(f"debug: got weapons_purchased 0x{weapons_purchased:0x}")
   all_weapons_mask = [i for i in WEAPONS if i["name"] == "all"][0]["p"]
   none_weapons_mask = [i for i in WEAPONS if i["name"] == "none"][0]["p"]
   # short-circuit if all
   if weapons_purchased & all_weapons_mask == all_weapons_mask:
      return "all", all_weapons_mask
   # short-circuit if none
   elif weapons_purchased | none_weapons_mask == 0:
      return "none", none_weapons_mask
   weapons_list = []
   weapons_mask = 0x0
   for i in WEAPONS:
      #result = weapons_purchased & i["p"]
      #print(f"debug: checking {i['p']:x} \"{i['name']}\" bitwise against {weapons_purchased:x}: {result:x}")
      if weapons_purchased & i["p"] and (i["name"] not in ["all","none"]):
         weapons_list.append(i["name"])
         weapons_mask += i["p"]
   if not silent:
      print(f"debug: currently have 0x{weapons_mask:04x} b{weapons_mask:016b}, {weapons_list}")
   return ','.join(weapons_list), weapons_mask

def get_level_status(data_object,profile_id,level):
   """ WORKHERE: not sure how to use this, but making sure my position in the data works. """
   data = _get_data_from_data_object(data_object)
   level_obj, message = get_level_info(level)
   if message != "valid" or level == -1:
      ferror(f"Unable to get level status for {level}.")
   print(f"Debug: got level_obj {level_obj} and message {message}")
   profile_level_status = data[PROFILE_START_POSITION[profile_id]+POS_LEVEL_START+(4*level_obj["pos_r"])]
   # it comes back as an int, but does it look better as a hex?
   return hex(profile_level_status)

def get_level_info(level):
   """ Returns dictionary of level from LEVELS, searching  by id or name. """
   try:
      # if it is an integer, make sure it shows up as one in the next check
      level = int(level)
   except:
      pass
   if type(level) == str:
      try:
         level = [i for i in LEVELS if i["name"].lower() == level.lower()][0]
         return level, "valid"
      except:
         return -1, f"cannot find level {level}"
   elif type(level) == int:
      if level in range(0,len(LEVELS)+1):
         try:
            level = [i for i in LEVELS if i["id"] == level][0]
            return level, "valid"
         except:
            return -1, f"cannot find level by id {level}"
      else:
         return -1, f"invalid level index {level}; use 0-{len(LEVELS)}"
   return -1, f"invalid way to reference level: [{type(level)}]. Use id or name."

def get_name(data_object,profile_id):
   """ Print the name of the profile_id. """
   data = _get_data_from_data_object(data_object)
   name_array = struct.unpack_from('<8i',data,PROFILE_START_POSITION[profile_id]+POS_NAME)
   #print(f"debug: name_array: {name_array}, type {type(name_array)}")
   name_str = ""
   for i in name_array:
      name_str += NAME_CHARS[i]
   return name_str

def set_name(data_object,profile_id,new_name):
   """ Set the name of the profile_id. Note that this allows you to set a name for a profile that is not in use, which preseeds the name when you select to make a profile in this slot in-game. """
   data = _get_data_from_data_object(data_object)
   # convert new_name to tuple of integer values
   name_array = []
   for i in new_name.upper():
      try:
         j = NAME_CHARS.index(i)
      except ValueError:
         return -1, f"Invalid characters in requested name {new_name}."
      name_array.append(j)
   # right-pad with spaces
   while len(name_array) < 8:
      name_array.append(0)
   if len(name_array) > 8 or len(name_array) < 0:
      return -1, f"Invalid length {len(name_array)} of name {new_name}."
   x = 0
   # How I was doing this before:
   #data_bytearray = bytearray(data)
   #for i in name_array:
   #   struct.pack_into('<1i',data_bytearray,PROFILE_START_POSITION[profile_id]+POS_NAME+(4*x),i)
   #   x = x + 1
   #data = bytes(data_bytearray)
   data = srb_pack('<8i',data,PROFILE_START_POSITION[profile_id]+POS_NAME+(4*x),*name_array)
   #print(f"debug: after setting name to {new_name}, we checked and got {get_name(data,profile_id)}")
   return data, ""

def get_profile_in_use(data_object,profile_id):
   """ Print if the profile_id is in use. """
   data = _get_data_from_data_object(data_object)
   # it always comes as a tuple
   in_use = struct.unpack_from('<1?',data,PROFILE_START_POSITION[profile_id]+POS_PROFILE_IN_USE)[0]
   #print(f"debug: in_use: {in_use}, type {type(in_use)}")
   return in_use

def get_tutorial_completed(data_object,profile_id):
   """ Print if the profile has completed the tutorial. """
   data = _get_data_from_data_object(data_object)
   # it always comes as a tuple
   tutorial_completed = struct.unpack_from('<1?',data,PROFILE_START_POSITION[profile_id]+POS_TUTORIAL_COMPLETED)[0]
   return tutorial_completed

def set_tutorial_completed(data_object,profile_id,completed):
   """ Set tutorial-completed for given profile. """
   data = _get_data_from_data_object(data_object)
   data = srb_pack('<?',data,PROFILE_START_POSITION[profile_id]+POS_TUTORIAL_COMPLETED,bool(completed))
   return data, get_tutorial_completed(data,profile_id)

def set_purchased_weapons(data_object,profile_id,action,weapons_list):
   """ For the given profile, take action on weapons_list, where action is in ["add","remove"] from the player. """
   data = _get_data_from_data_object(data_object)
   if action not in ["add","remove"]:
      return -1, f"Failed: can only [\"add\",\"remove\"] purchased weapons"
   # validate weapons_list
   action_mask = 0x0
   for w in weapons_list:
      weapon, message = get_weapon_info(w)
      if message != "":
         return -1, f"unable to {action} weapons because {message} on weapon {w}"
      action_mask = action_mask | weapon["p"]
   cur_weapons, cur_mask = get_purchased_weapons(data,profile_id)
   #print(f"debug: action_mask(type {type(action_mask)})={action_mask}")
   #print(f"debug: cur_mask(type {type(cur_mask)})={cur_mask}")
   #print(f"debug: need to {action}-combine {cur_mask:016b} and {action_mask:016b}")
   if action == "add":
      final_mask = cur_mask | action_mask
   elif action == "remove":
      final_mask = cur_mask & ~action_mask
   if final_mask != cur_mask:
      print(f"debug: beginning  0x{cur_mask:04x}, b{cur_mask:016b} {cur_weapons}")
      print(   f"debug: {action:6s}     0x{action_mask:04x}, b{action_mask:016b} {','.join(weapons_list)}")
   data = srb_pack('<1i',data,PROFILE_START_POSITION[profile_id]+POS_BYTES_WEAPONS_PURCHASED, final_mask)
   print(   f"debug: final      0x{final_mask:04x}, b{final_mask:016b} {get_purchased_weapons(data,profile_id,silent=True)[0]}")
   # if we make it to the end
   return data, ""

def srb_pack(format_str, data, offset, *new_contents):
   """ Helper function that accepts data as bytes, instead of requiring bytesarray. """
   data_bytearray = bytearray(data)
   struct.pack_into(format_str,data_bytearray,offset,*new_contents)
   return bytes(data_bytearray)

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"debug: Got csum={csum}")
   # pad any odd-digit-count hex value with a leading 0 for bytearray.fromhex() which is not very smart.
   csum = ("0" if len(str(csum)) % 2 else "") + str(csum)[2:]
   csum = bytearray.fromhex(csum)
   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