add get/add/remove purchased planes

2 files changed, 85 insertions, 16 deletions

diff --git a/srb.py b/srb.py
index f97e01b..81c8b5f 100755
--- a/srb.py
+++ b/srb.py
@@ -9,8 +9,8 @@
 #    delsum from github
 #    bgconf.py
 
-# WORKHERE: right now, this only just corrects any file. in the future, offer options like --take-plane all or --give-plane marcie
-# or --give-plane marcie or --set-rank-level 1,general --set-rank-level 1,none --set-rank-level all,general
+# WORKHERE:
+# or --set-rank-level 1,general --set-rank-level 1,none --set-rank-level all,general
 # or --give-balloons-level 10 --take-balloons-level 11
 # or --reset-level 5 (which zeros out all info about completion of that level)
 
@@ -39,9 +39,12 @@ parser.add_argument("--get-weapon",action="store_true",help="Print currently equ
 #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("--get-purchased-weapons",action="store_true",help="Print currently purchased weapons 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, remove (un-buy) these purchased weapons. Can be used multiple times.")
+parser.add_argument("--get-purchased-planes",action="store_true",help="Print currently purchased planes for profile.")
+parser.add_argument("--add-purchased-planes",action="append",help="For profile, add these purchased planes. Can be used multiple times.")
+parser.add_argument("--remove-purchased-planes",action="append",help="For profile, remove (un-buy) these purchased planes. Can be used multiple times.")
 parser.add_argument("--get-level",help="Print status for this level for profile.")
 parser.add_argument("--get-levelset",help="Print status for this levelset for profile.")
 parser.add_argument("--get-name",action="store_true",help="Print name for profile.")
@@ -73,7 +76,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 or args.get_purchased_weapons or args.get_tutorial_completed or args.add_purchased_weapons or args.remove_purchased_weapons or args.get_health or args.get_stunt or args.get_gun or args.set_health or args.set_stunt or args.set_gun or args.get_levelset):
+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 or args.get_health or args.get_stunt or args.get_gun or args.set_health or args.set_stunt or args.set_gun or args.get_levelset or args.get_purchased_planes or args.add_purchased_planes or args.remove_purchased_planes):
    ferror("Warning: Cannot perform most actions without --profile. Not all tasks may run.")
 else:
    if args.get_money:
@@ -165,6 +168,20 @@ else:
          print(f"Failed to set gun to {args.set_gun} because {message}")
       else:
          srb_lib.write_file(args.file,0,data)
+   if args.get_purchased_planes:
+      print(f"Profile {profile_id} has planes {srb_lib.get_purchased_planes(args.file,profile_id)}")
+   if args.add_purchased_planes:
+      data, message = srb_lib.set_purchased_planes(args.file,profile_id,"add",args.add_purchased_planes)
+      if (type(data) == int and data == -1) or message != "":
+         ferror(f"Failed to add purchased planes {args.add_purchased_planes} because {message}")
+      else:
+         srb_lib.write_file(args.file,0,data)
+   if args.remove_purchased_planes:
+      data, message = srb_lib.set_purchased_planes(args.file,profile_id,"remove",args.remove_purchased_planes)
+      if (type(data) == int and data == -1) or message != "":
+         ferror(f"Failed to add purchased planes {args.remove_purchased_planes} because {message}")
+      else:
+         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 e8adab4..253a3be 100644
--- a/srb_lib.py
+++ b/srb_lib.py
@@ -20,7 +20,7 @@
 # Dependencies:
 
 import sys, struct
-srb_lib_version = "20240313a"
+srb_lib_version = "20240318a"
 
 # 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.
 # "Z<dddddddd" is the start of a profile.
@@ -39,16 +39,6 @@ POS_NAME = 0x294
 
 NAME_CHARS = " ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890!?"
 
-# absolute 0x29C, so pos 0x288 is where the purchased planes are stored. '<1i' unpack
-# WORKHERE: this is the bitmask to use for giving purchased planes
-# marcie      0x0100 b00000100000000
-# sally       0x0020 b00000000100000
-# rerun       0x0800 b00100000000000
-# pigpen      0x0200 b00001000000000
-# woodstock   0x1000 b01000000000000
-# baron       0x2000 b10000000000000
-# all         0x3B20 b11101100100000
-
 POS_BYTES_WEAPONS_PURCHASED = 0x028C
 # Woodstock Missile and Bottle Rockets are always available to use.
 # for equipped-weapon
@@ -74,6 +64,18 @@ WEAPONS = [
    {"id":0xF,"e":False,"p":0x6F60,"name":"all"}, # custom for this library
 ]
 
+POS_PLANES_PURCHASED = 0x288
+PLANES = [
+   {"p":0x0000,"name":"none"},
+   {"p":0x0100,"name":"marcie"},
+   {"p":0x0020,"name":"sally"},
+   {"p":0x0800,"name":"rerun"},
+   {"p":0x0200,"name":"pigpen"},
+   {"p":0x1000,"name":"woodstock"},
+   {"p":0x2000,"name":"baron"},
+   {"p":0x3B20,"name":"all"},
+]
+
 # c = balloon color
 # l = level count
 # le = letters, calculated
@@ -95,7 +97,6 @@ POS_LEVEL_BALLOONS_MULTIPLIER_LEVELSET = 0x238
 POS_LEVEL_BALLOONS_MULTIPLIER_LEVEL = 0x50
 POS_LEVELSET_COMPLETED_LETTERS_COUNT = 0x3A0
 POS_LEVEL_LETTERS = 0x3E8
-#POS_LEVEL_LETTERS_MULTIPLIER_LEVELSET = 0x1C
 POS_LEVEL_LETTERS_MULTIPLIER_LEVELSET = 7
 POS_LEVELSET_COMPLETED_MISSIONS_MASK = 0x2B8
 
@@ -604,6 +605,57 @@ def set_purchased_weapons(data_object,profile_id,action,weapons_list):
    # if we make it to the end
    return data, ""
 
+def get_purchased_planes(data_object,profile_id, silent=False):
+   """ List which planes are purchased already. """
+   data = _get_data_from_data_object(data_object)
+   planes_purchased = struct.unpack_from('<1i',data,PROFILE_START_POSITION[profile_id]+POS_PLANES_PURCHASED)[0]
+   all_planes_mask = [i for i in PLANES if i["name"] == "all"][0]["p"]
+   none_planes_mask = [i for i in PLANES if i["name"] == "none"][0]["p"]
+   # short-circuit if all
+   if planes_purchased & all_planes_mask == all_planes_mask:
+      return "all", all_planes_mask
+   # short-circuit if none
+   elif planes_purchased | none_planes_mask == 0:
+      return "none", none_planes_mask
+   planes_list = []
+   planes_mask = 0x0
+   for i in PLANES:
+      if planes_purchased & i["p"] and (i["name"] not in ["all","none"]):
+         planes_list.append(i["name"])
+         planes_mask += i["p"]
+   if not silent:
+      print(f"debug: currently have 0x{planes_mask:04x} b{planes_mask:016b}, {planes_list}")
+   return ','.join(planes_list), planes_mask
+
+def set_purchased_planes(data_object,profile_id,action,planes_list):
+   """ For the given profile, take action on planes_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 planes"
+   # validate planes_list
+   action_mask = 0x0
+   for p in planes_list:
+      try:
+         plane = [i for i in PLANES if i["name"] == p][0]
+      except:
+         return -1, f"unable to {action} planes because {message} on plane {p}"
+      action_mask = action_mask | plane["p"]
+   cur_planes, cur_mask = get_purchased_planes(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_planes}")
+      print(   f"debug: {action:6s}     0x{action_mask:04x}, b{action_mask:016b} {','.join(planes_list)}")
+   data = srb_pack('<1i',data,PROFILE_START_POSITION[profile_id]+POS_PLANES_PURCHASED, final_mask)
+   print(   f"debug: final      0x{final_mask:04x}, b{final_mask:016b} {get_purchased_planes(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)