diff --git a/s1ap_reiniger.py b/s1ap_reiniger.py
index f121134..b56e45f 100755
--- a/s1ap_reiniger.py
+++ b/s1ap_reiniger.py
@@ -25,6 +25,24 @@ from scapy.all import SCTP
 
 from pycrate_asn1dir import S1AP
 from pycrate_mobile import NAS
+from pycrate_core.charpy import Charpy
+
+IMSI_REPLACEMENT_BYTES = b'\x29\x26\x24' + b'\x00' * 5
+
+def get_key_or_none(elem, k):
+	try:
+		id = elem[k]
+		return id
+	except ValueError:
+		return None
+
+def has_key(elem, k):
+	try:
+		id = elem[k]
+		return True
+	except:
+		return False
+
 
 def handle_nas_pdu(pdu, dl, regen = False):
 	log.debug("Processing %s NAS PDU: %s" % ("Downlink" if dl else "Uplink", pdu.hex()))
@@ -32,19 +50,34 @@ def handle_nas_pdu(pdu, dl, regen = False):
 	if code:
 		log.error("Failed to parse NAS payload")
 		return None
-	
+	#print(msg.CLASS)
+
+	if has_key(msg, 'ID'):
+		id = msg['ID'][1]
+		#for k in id:
+		#	print("--- %s, %s" % (k, k.show()))
+		id_type = id['Type'].get_val()
+		if id_type == 1: # IMSI
+			id.from_bytes(Charpy(IMSI_REPLACEMENT_BYTES))
+		else:
+			raise FooErr
+		print("+++ %s" % (id))
+		regen = True
+
 	# Try to find EPSID (may contain IMSI)
 	# TODO: also patch IMEI / IMEISV
-	try:
+	if has_key(msg, 'EPSID'):
 		epsid = msg['EPSID'][1]
 		# Check if EPSID contains exactly IMSI
-		if epsid[2].get_val() == 1:
+		id_type = epsid['Type'].get_val()
+		print("ID type: %d" % (id_type))
+		if id_type == 1:
 			log.info("Cleaning %s" % epsid.repr())
 			# 262420000000000, Vodafone GmbH, Germany
-			epsid.from_bytes('\x29\x26\x24' + '\x00' * 5)
-			regen = True
-	except:
-		pass
+			epsid.from_bytes(Charpy(IMSI_REPLACEMENT_BYTES))
+		else:
+			raise FooErr
+		regen = True
 
 	return msg.to_bytes() if regen else pdu