Add -r / --re-encrypt

libray/core.py

@@ -102,7 +102,7 @@ def download_ird(ird_name):
   """Download an .ird from GET_IRD_NET_LOC"""
 
   # Check if file already exists and skip if it does
-  if os.path.exists(ird_name):
+  if os.path.exists(ird_name): # TODO: might want to check that the file is valid first, could do a HEAD agains the url
     return
 
   ird_link = GET_IRD_NET_LOC + ird_name
@@ -119,7 +119,7 @@ def ird_by_game_id(game_id):
   try:
     r = requests.get(ALL_IRD_NET_LOC, headers = {'User-Agent': 'Anonymous (You)' }, timeout=5)
   except requests.exceptions.ReadTimeout:
-    core.error('Server timed out, fix your connection or manually specify a key/ird.')
+    error('Server timed out, fix your connection or manually specify a key/ird.')
   soup = BeautifulSoup(r.text, "html.parser")
 
   ird_name = False
@@ -140,7 +140,7 @@ def ird_by_game_id(game_id):
 
 
 def decrypt(args):
-  """Try to decrypt a given .iso using relevant .ird using args from argparse
+  """Try to decrypt a given .iso using relevant .ird or encryption key from argparse
 
   If no .ird is given this will try to automatically download an .ird file with the encryption/decryption key for the given game .iso
   """
@@ -150,4 +150,12 @@ def decrypt(args):
   input_iso.decrypt(args)
 
 
+def encrypt(args):
+  """Try to re-encrypt a decrypted .iso using relevant .ird or encryption key from argparse
 
+  If no .ird is given this will try to automatically download an .ird file with the encryption/decryption key for the given game .iso
+  """
+
+  input_iso = iso.ISO(args)
+
+  input_iso.encrypt(args)

libray/iso.py

@@ -180,10 +180,67 @@ class ISO:
           print('Decryption complete.')
 
 
+  def encrypt(self, args):
+    """Encrypt self using args from argparse."""
+
+    if not args.quiet:
+      print('Re-encrypting with disc key: %s' % self.disc_key.hex())
+
+    with open(args.iso, 'rb') as input_iso:
+
+      if not args.output:
+        output_name = '%s_e.iso' % self.game_id
+      else:
+        output_name = args.output
+
+      with open(output_name, 'wb') as output_iso:
+
+        if not args.quiet:
+          pbar = tqdm(total= (self.size // 2048) )
+
+        for region in self.regions:
+          input_iso.seek(region['start'])
+
+          # Unencrypted region, just copy it
+          if not region['enc']:
+            while input_iso.tell() < region['end']:
+              data = input_iso.read(core.SECTOR)
+              if not data and not args.quiet:
+                core.warning('Trying to read past the end of the file')
+                break
+              output_iso.write(data)
+
+              if not args.quiet:
+                pbar.update(1)
+            continue
+          # Decrypted region, re-encrypt it
+          else:
+            while input_iso.tell() < region['end']:
+              num = input_iso.tell() // 2048
+              iv = bytearray([0 for i in range(0,16)])
+              for j in range(0,16):
+                iv[16 - j - 1] = (num & 0xFF)
+                num >>= 8
+
+              data = input_iso.read(core.SECTOR)
+              if not data and not args.quiet:
+                core.warning('Trying to read past the end of the file')
+                break
+
+              cipher = AES.new(self.disc_key, AES.MODE_CBC, bytes(iv))
+              encrypted = cipher.encrypt(data)
+
+              output_iso.write(encrypted)
+
+              if not args.quiet:
+                pbar.update(1)
+
+
   def print_info(self):
     # TODO: This could probably have been a __str__? Who cares?
     """Print some info about the ISO."""
-
+    print('Game ID: %s' % self.game_id)
+    print('Key: %s' % self.disc_key.hex())
     print('Info from ISO:')
     print('Regions: %s' % self.number_of_regions)
     for i, region in enumerate(self.regions):

libray/libray.py

@@ -43,7 +43,15 @@ if __name__ == '__main__':
   optional.add_argument('-d', '--decryption-key', dest='decryption_key', type=str, help='Manually specify key', default='')
   optional.add_argument('-v', '--verbose', dest='verbose', help='Increase verbosity', action='count')
   optional.add_argument('-q', '--quiet', dest='quiet', help='Quiet mode, only prints on error', action='store_true')
+  # -e is reserved for "extract" so re-encrypt is "-r"
+  optional.add_argument('-r', '--re-encrypt', dest='reencrypt', help='Re-encrypt .iso', action='store_true')
 
   args = parser.parse_args()
 
-  core.decrypt(args)
+  if args.reencrypt:
+
+    core.encrypt(args)
+
+  else:
+
+    core.decrypt(args)