Merge pull request #19 from rtpt-romankarwacik/handle_deserialization_errors

Handle deseralization errors

Author: ShutdownRepo

pywhisker/pywhisker.py

@@ -283,6 +283,9 @@ def info(self, device_id):
             device_id_in_current_values = False
             for dn_binary_value in results['raw_attributes']['msDS-KeyCredentialLink']:
                 keyCredential = KeyCredential.fromDNWithBinary(DNWithBinary.fromRawDNWithBinary(dn_binary_value))
+                if keyCredential.DeviceId is None:
+                    logger.warning("Failed to parse DeviceId for keyCredential: %s" % (str(dn_binary_value)))
+                    continue
                 if keyCredential.DeviceId.toFormatD() == device_id:
                     logger.success("Found device Id")
                     keyCredential.show()
@@ -319,7 +322,12 @@ def list(self):
                 logger.info("Listing devices for %s" % self.target_samname)
                 for dn_binary_value in results['raw_attributes']['msDS-KeyCredentialLink']:
                     keyCredential = KeyCredential.fromDNWithBinary(DNWithBinary.fromRawDNWithBinary(dn_binary_value))
-                    logger.info("DeviceID: %s | Creation Time (UTC): %s" % (keyCredential.DeviceId.toFormatD(), keyCredential.CreationTime))
+                    if keyCredential.DeviceId is None:
+                        logger.warning("Failed to parse DeviceId for keyCredential: %s" % (str(dn_binary_value)))
+                        logger.warning("DeviceID: %s | Creation Time (UTC): %s" % (keyCredential.DeviceId, keyCredential.CreationTime))
+                    else:
+                        logger.info("DeviceID: %s | Creation Time (UTC): %s" % (keyCredential.DeviceId.toFormatD(), keyCredential.CreationTime))
+
         except IndexError:
             logger.info('Attribute msDS-KeyCredentialLink does not exist')
         return
@@ -466,6 +474,9 @@ def remove(self, device_id):
             device_id_in_current_values = False
             for dn_binary_value in results['raw_attributes']['msDS-KeyCredentialLink']:
                 keyCredential = KeyCredential.fromDNWithBinary(DNWithBinary.fromRawDNWithBinary(dn_binary_value))
+                if keyCredential.DeviceId is None:
+                    logger.warning("Failed to parse DeviceId for keyCredential: %s" % (str(dn_binary_value)))
+                    continue
                 if keyCredential.DeviceId.toFormatD() == device_id:
                     logger.info("Found value to remove")
                     device_id_in_current_values = True
@@ -553,7 +564,10 @@ def importFromJSON(self, filename):
                 with open(filename, "r") as f:
                     data = json.load(f)
                     for kcjson in data["keyCredentials"]:
-                        keyCredentials.append(KeyCredential.fromDict(kcjson).toDNWithBinary().toString())
+                        if type(kcjson) == dict:
+                            keyCredentials.append(KeyCredential.fromDict(kcjson).toDNWithBinary().toString())
+                        elif type(kcjson) == str:
+                            keyCredentials.append(kcjson)
             logger.info("Modifying the msDS-KeyCredentialLink attribute of %s" % self.target_samname)
             self.ldap_session.modify(self.target_dn, {'msDS-KeyCredentialLink': [ldap3.MODIFY_REPLACE, keyCredentials]})
             if self.ldap_session.result['result'] == 0:
@@ -599,7 +613,11 @@ def exportToJSON(self, filename):
             keyCredentialsJSON = {"keyCredentials":[]}
             for dn_binary_value in results['raw_attributes']['msDS-KeyCredentialLink']:
                 keyCredential = KeyCredential.fromDNWithBinary(DNWithBinary.fromRawDNWithBinary(dn_binary_value))
-                keyCredentialsJSON["keyCredentials"].append(keyCredential.toDict())
+                try:
+                    keyCredentialsJSON["keyCredentials"].append(keyCredential.toDict())
+                except Exception as e:
+                    logger.warning(f"Failed to serialize keyCredential, error: %s, saving the raw keyCredential instead, i.e.: %s" % (str(e), dn_binary_value.decode()))
+                    keyCredentialsJSON["keyCredentials"].append(dn_binary_value.decode())
             with open(filename, "w") as f:
                 f.write(json.dumps(keyCredentialsJSON, indent=4))
             logger.success("Saved JSON dump at path: %s" % filename)