add get/set name

2 files changed, 67 insertions, 5 deletions

diff --git a/srb.py b/srb.py
index 800d1c5..9c6a941 100755
--- a/srb.py
+++ b/srb.py
@@ -27,6 +27,7 @@ parser = argparse.ArgumentParser(description="Cli tool for manipulating savegame
    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))}
+NAME_CHARS include "{srb_lib.NAME_CHARS}"
 """)
 parser.add_argument("-V","--version",action="version",version="%(prog)s " + srb_lib.srb_lib_version)
 parser.add_argument("-d","--debug",nargs='?',default=0,type=int,choices=range(0,11),help="Set debug level")
@@ -37,8 +38,11 @@ parser.add_argument("--get-weapon",action="store_true",help="Print currently equ
 # 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="Print currently equipped weapon for profile.")
+parser.add_argument("--set-weapon",help="Set currently equipped weapon for profile.")
 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("--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()
@@ -57,7 +61,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):
+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):
    ferror("Warning: Cannot perform most actions without --profile. Not all tasks may run.")
 else:
    if args.get_money:
@@ -84,6 +88,16 @@ else:
          srb_lib.write_file(args.file,0,data)
    if args.get_level:
       print(f"Profile {profile_id} has level {args.get_level} status {srb_lib.get_level_status(args.file,profile_id,args.get_level)}")
+   if args.get_name:
+      print(f"Profile {profile_id} has name {srb_lib.get_name(args.file,profile_id)}")
+   if args.set_name:
+      data, message = srb_lib.set_name(args.file, profile_id, args.set_name)
+      if (type(data) == int and data == -1) or message != "":
+         ferror(f"Failed to set profile {profile_id} name to {args.set_name} because {message}")
+      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.checksum:
    f = args.file
diff --git a/srb_lib.py b/srb_lib.py
index dc69614..1b13175 100644
--- a/srb_lib.py
+++ b/srb_lib.py
@@ -18,16 +18,19 @@
 #    winetricks vd=1024x768
 #    winetricks vd=off
 # Dependencies:
-#    python3-crcmod
 
 import sys, struct
-srb_lib_version = "20240309a"
+srb_lib_version = "20240310a"
 
 # 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!?"
 
 # stinger purchased is relative 0x28C, value 0x40 (b01000000) or 64
 # water-balloon-gun is relative 0x2D, value 0x0001
@@ -190,6 +193,7 @@ def crc_poly(data, n, poly, crc=0, ref_in=False, ref_out=False, xor_out=0):
 
 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:
@@ -268,7 +272,7 @@ def get_level_status(data_object,profile_id,level):
    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:
@@ -293,6 +297,50 @@ def get_level_info(level):
          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)
+   data_bytearray = bytearray(data)
+   # 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
+   # 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)
+   #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 calculate_checksum(data):
    """ Return the 4-byte checksum used by the game for the provided data. """
    # aka # CRC-32/BZIP2