add get/add/remove purchased planes

1 files changed, 64 insertions, 12 deletions

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)