add get/set plane stats

2 files changed, 79 insertions, 2 deletions

diff --git a/srb.py b/srb.py
index 658c272..01974f8 100755
--- a/srb.py
+++ b/srb.py
@@ -16,7 +16,7 @@
 # make chart of level numbers, which ones have balloons
 # or --reset-level 5 (which zeros out all info about completion of that level)
 # --set-plane-health 1 (through 4)
-# --set-plane-weapon 1 (through 4)
+# --set-plane-weapon 1 (through 5)
 # --set-plane-stunt 1 (through 4)
 # --set-profile-name "asdbdf"
 
@@ -54,6 +54,12 @@ 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("--get-health",action="store_true",help="Print stat health for profile.")
+parser.add_argument("--set-health",type=int,choices=range(1,5),help="Set stat health for profile.")
+parser.add_argument("--get-stunt",action="store_true",help="Print stat stunt for profile.")
+parser.add_argument("--set-stunt",type=int,choices=range(1,5),help="Set stat stunt for profile.")
+parser.add_argument("--get-gun",action="store_true",help="Print stat gun for profile.")
+parser.add_argument("--set-gun",type=int,choices=range(1,6),help="Set stat gun 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()
@@ -72,7 +78,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):
+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):
    ferror("Warning: Cannot perform most actions without --profile. Not all tasks may run.")
 else:
    if args.get_money:
@@ -129,6 +135,39 @@ else:
       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.get_health:
+      stat, message = srb_lib.get_plane_stat(args.file,profile_id,"health")
+      # error is printed in the function, so if not any error, run
+      if not (stat == -1 or message != ""):
+         print(f"Profile {profile_id} has health level {stat}")
+   if args.set_health:
+      data, message = srb_lib.set_plane_stat(args.file,profile_id,"health",args.set_health)
+      if data == -1 or message != "":
+         print(f"Failed to set health to {args.set_health} because {message}")
+      else:
+         srb_lib.write_file(args.file,0,data)
+   if args.get_stunt:
+      stat, message = srb_lib.get_plane_stat(args.file,profile_id,"stunt")
+      # error is printed in the function, so if not any error, run
+      if not (stat == -1 or message != ""):
+         print(f"Profile {profile_id} has stunt level {stat}")
+   if args.set_stunt:
+      data, message = srb_lib.set_plane_stat(args.file,profile_id,"stunt",args.set_stunt)
+      if data == -1 or message != "":
+         print(f"Failed to set stunt to {args.set_stunt} because {message}")
+      else:
+         srb_lib.write_file(args.file,0,data)
+   if args.get_gun:
+      stat, message = srb_lib.get_plane_stat(args.file,profile_id,"gun")
+      # error is printed in the function, so if not any error, run
+      if not (stat == -1 or message != ""):
+         print(f"Profile {profile_id} has gun level {stat}")
+   if args.set_gun:
+      data, message = srb_lib.set_plane_stat(args.file,profile_id,"gun",args.set_gun)
+      if data == -1 or message != "":
+         print(f"Failed to set gun to {args.set_gun} 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 32a4f2e..192fc22 100644
--- a/srb_lib.py
+++ b/srb_lib.py
@@ -26,6 +26,9 @@ srb_lib_version = "20240311a"
 # money is 0x270 bytes after the "Z<dddddddd" start.
 PROFILE_START_POSITION = [0, 0x10, 0x142C, 0x2848]
 POS_MONEY = 0x270
+POS_HEALTH = 0x274 # 0-3 is available levels
+POS_STUNT = 0x280 # 0-3 is available levels
+POS_GUN = 0x27C # 0-3 is available levels
 POS_EQUIPPED_WEAPON = 0x284
 POS_PROFILE_IN_USE = 0x290
 POS_NAME = 0x294
@@ -71,6 +74,7 @@ LEVELSETS = [
 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.
+# WORKHERE 2024-03-11-2 19:29 porifle 3 replay mission 0, get the letter and maybe a few more balloons?
 LEVELS = [
    {"id":0, "pos_r":0, "setid":0,"name":"Defend Island"},
    {"id":1, "pos_r":1, "setid":0,"name":"Recover Plans"},
@@ -313,6 +317,7 @@ def get_level_status(data_object,profile_id,level):
       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"])]
+   print(f"Debug: got level status {profile_level_status:x} at pos 0x{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)
 
@@ -398,6 +403,39 @@ def set_tutorial_completed(data_object,profile_id,completed):
    data = srb_pack('<?',data,PROFILE_START_POSITION[profile_id]+POS_TUTORIAL_COMPLETED,bool(completed))
    return data, get_tutorial_completed(data,profile_id)
 
+def get_plane_stat(data_object,profile_id,stat):
+   """ Get health/stunt/gun statistic for profile. """
+   data = _get_data_from_data_object(data_object)
+   stat = stat.lower()
+   if stat not in ["health","stunt","gun"]:
+      return -1, f"Available plane stats are health,stunt,gun."
+   if "health" == stat:
+      thispos = POS_HEALTH
+   elif "stunt" == stat:
+      thispos = POS_STUNT
+   else: # gun
+      thispos = POS_GUN
+   output = struct.unpack_from('<1i',data,PROFILE_START_POSITION[profile_id]+thispos)[0]+1
+   #print(f"debug: at offset {PROFILE_START_POSITION[profile_id]+thispos:x} got {output:1b}")
+   return output, ""
+
+def set_plane_stat(data_object,profile_id,stat,value):
+   """ Set health/stunt/gun statistic for profile. """
+   data = _get_data_from_data_object(data_object)
+   stat = stat.lower()
+   if stat not in ["health","stunt","gun"]:
+      return -1, f"Available plane stats are health,stunt,gun."
+   if value not in range(1,6) or (value == 5 and stat != "gun"):
+      return -1, f"Invalid value {value} for stat."
+   if "health" == stat:
+      thispos = POS_HEALTH
+   elif "stunt" == stat:
+      thispos = POS_STUNT
+   else: # gun
+      thispos = POS_GUN
+   data = srb_pack('<1i',data,PROFILE_START_POSITION[profile_id]+thispos,value-1)
+   return data, ""
+
 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)