get/set purchased weapons

2 files changed, 177 insertions, 53 deletions

diff --git a/srb.py b/srb.py
index 9c6a941..658c272 100755
--- a/srb.py
+++ b/srb.py
@@ -23,10 +23,16 @@
 import srb_lib, argparse, sys
 from srb_lib import ferror, debuglev
 
+HELP_WEAPONS = ""
+
+for w in srb_lib.WEAPONS:
+   HELP_WEAPONS += str(w["id"]) + "," + w["name"] + ','
+HELP_WEAPONS = HELP_WEAPONS.rstrip(",")
+
 parser = argparse.ArgumentParser(description="Cli tool for manipulating savegame files for srb.exe",
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog=f"""LEVELS include {[i['name'] for i in srb_lib.LEVELS]+list(range(0,len(srb_lib.LEVELS)))}
-WEAPONS include {[i for i in srb_lib.WEAPONS if i != "undefined"]+list(range(0,16))}
+WEAPONS include {HELP_WEAPONS}
 NAME_CHARS include "{srb_lib.NAME_CHARS}"
 """)
 parser.add_argument("-V","--version",action="version",version="%(prog)s " + srb_lib.srb_lib_version)
@@ -35,14 +41,19 @@ parser.add_argument("--profile",type=int,choices=range(1,4),help="Profile in use
 parser.add_argument("--get-money",action="store_true",help="Print current money for profile.")
 parser.add_argument("--set-money",type=int,help="Set money for profile.")
 parser.add_argument("--get-weapon",action="store_true",help="Print currently equipped weapon for profile.")
-# choices seems to be too strict here for the numbers. We can live with just the 
+# choices seems to be too strict here for the numbers. We can live with just the
 #choices=[i for i in srb_lib.WEAPONS if i != "undefined"]+list(range(0,16))
 #parser.add_argument("--set-weapon",choices=[i for i in srb_lib.WEAPONS if i != "undefined"],help="Print currently equipped weapon for profile.")
 parser.add_argument("--set-weapon",help="Set currently equipped weapon for profile.")
+parser.add_argument("--get-purchased-weapons",action="store_true",help="Print currently purchased weapon for profile.")
+parser.add_argument("--add-purchased-weapons",action="append",help="For profile, add these purchased weapons. Can be used multiple times.")
+parser.add_argument("--remove-purchased-weapons",action="append",help="For profile, add these purchased weapons. Can be used multiple times.")
 parser.add_argument("--get-level",help="Print status for this level for profile.")
 parser.add_argument("--get-name",action="store_true",help="Print name for profile.")
 parser.add_argument("--set-name",help="Set name for profile.")
 parser.add_argument("--get-profile-in-use",action="store_true",help="Print if profile is in use.")
+parser.add_argument("--get-tutorial-completed",action="store_true",help="Print if profile has completed the tutorial.")
+parser.add_argument("--set-tutorial-completed",choices=["True","False"],help="Set tutorial-completed for profile.")
 parser.add_argument("--checksum",action=argparse.BooleanOptionalAction,default=True,help="Correct checksum. Default is to do this. It happens at the end of everything else.")
 parser.add_argument("file",default="Profile 1.sav")
 args = parser.parse_args()
@@ -61,7 +72,7 @@ profile_id = args.profile
 #print(f"profile_id={profile_id}")
 
 # WORKHERE: new actions that need --profile must be added here.
-if not profile_id and (args.get_money or args.set_money or args.get_weapon or args.set_weapon or args.get_level or args.get_name or args.set_name or args.get_profile_in_use):
+if not profile_id and (args.get_money or args.set_money or args.get_weapon or args.set_weapon or args.get_level or args.get_name or args.set_name or args.get_profile_in_use or args.get_purchased_weapons or args.get_tutorial_completed or args.add_purchased_weapons or args.remove_purchased_weapons):
    ferror("Warning: Cannot perform most actions without --profile. Not all tasks may run.")
 else:
    if args.get_money:
@@ -98,6 +109,26 @@ else:
          srb_lib.write_file(args.file,0,data)
    if args.get_profile_in_use:
       print(f"Profile {profile_id} in use is {srb_lib.get_profile_in_use(args.file,profile_id)}")
+   if args.get_purchased_weapons:
+      print(f"Profile {profile_id} has weapons {srb_lib.get_purchased_weapons(args.file,profile_id)}")
+   if args.add_purchased_weapons:
+      data, message = srb_lib.set_purchased_weapons(args.file,profile_id,"add",args.add_purchased_weapons)
+      if (type(data) == int and data == -1) or message != "":
+         ferror(f"Failed to add purchased weapons {args.add_purchased_weapons} because {message}")
+      else:
+         srb_lib.write_file(args.file,0,data)
+   if args.remove_purchased_weapons:
+      data, message = srb_lib.set_purchased_weapons(args.file,profile_id,"remove",args.remove_purchased_weapons)
+      if (type(data) == int and data == -1) or message != "":
+         ferror(f"Failed to add purchased weapons {args.remove_purchased_weapons} because {message}")
+      else:
+         srb_lib.write_file(args.file,0,data)
+   if args.get_tutorial_completed:
+      print(f"Profile {profile_id} completed-tutorial is {srb_lib.get_tutorial_completed(args.file,profile_id)}")
+   if args.set_tutorial_completed:
+      thisbool = False if args.set_tutorial_completed == "False" else True
+      data, newstatus = srb_lib.set_tutorial_completed(args.file,profile_id,thisbool)
+      srb_lib.write_file(args.file,0,data)
 
 if args.checksum:
    f = args.file
diff --git a/srb_lib.py b/srb_lib.py
index 1b13175..32a4f2e 100644
--- a/srb_lib.py
+++ b/srb_lib.py
@@ -20,7 +20,7 @@
 # Dependencies:
 
 import sys, struct
-srb_lib_version = "20240310a"
+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.
@@ -32,28 +32,29 @@ POS_NAME = 0x294
 
 NAME_CHARS = " ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890!?"
 
-# stinger purchased is relative 0x28C, value 0x40 (b01000000) or 64
-# water-balloon-gun is relative 0x2D, value 0x0001
-# stinger and potato gun is relative 0x28C, value 0x60, so potato gun is (b00100000) or 64
-#POS_STINGER = 0x028C
-
+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 = [
-   "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
+   {"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.
@@ -107,6 +108,8 @@ LEVEL_STATUSES = {
    "64": "all-balloons,all-letters,general"
 }
 
+POS_TUTORIAL_COMPLETED = 0x2D0 # bool
+
 # DEFINE FUNCTIONS
 def ferror(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)
@@ -222,45 +225,85 @@ 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)
+   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 = data[PROFILE_START_POSITION[profile_id]+POS_EQUIPPED_WEAPON]
-   return WEAPONS[weapon]
+   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)
-   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}")
+   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
-   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)}")
+   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 index, and name of weapon, or -1 and "invalid" """
+   """
+   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_id = [index for index in range(len(WEAPONS)) if WEAPONS[index].lower() == weapon.lower()][0]
-         return weapon_id, WEAPONS[weapon_id]
+         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 -1, f"cannot find weapon {weapon}" # must be an incorrect name
+         return {}, f"cannot find weapon {weapon}" # must be an incorrect name
    elif type(weapon) == int:
       if weapon in range(0,16):
-         return weapon, WEAPONS[weapon]
+         weapon_obj = [i for i in WEAPONS if i["id"] == weapon][0]
+         return weapon_obj, ""
+         #return weapon, WEAPONS[weapon]
       else:
-         return -1, f"invalid index {weapon}; use 0-15" # must be <0 or >15
-   return -1, f"invalid way to reference weapon: [{type(weapon)}]. Use index or name."
+         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. """
@@ -310,7 +353,6 @@ def get_name(data_object,profile_id):
 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)
-   data_bytearray = bytearray(data)
    # convert new_name to tuple of integer values
    name_array = []
    for i in new_name.upper():
@@ -325,11 +367,13 @@ def set_name(data_object,profile_id,new_name):
    if len(name_array) > 8 or len(name_array) < 0:
       return -1, f"Invalid length {len(name_array)} of name {new_name}."
    x = 0
-   # for some reason I cannot figure out '<8i' with a tuple or list of ints. So just do it manually.
-   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)
+   # 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, ""
 
@@ -341,6 +385,53 @@ def get_profile_in_use(data_object,profile_id):
    #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
@@ -350,8 +441,10 @@ def calculate_checksum(data):
    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:])
+   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)