Fix code, add more features

AirTagCrypto/AirTagCrypto.py

@@ -1,6 +1,5 @@
 import base64
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
-from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat
 from cryptography.hazmat.primitives.asymmetric import ec
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import hashes
@@ -13,14 +12,30 @@ def __init__(self, private_key: str = None):
         else:
             self._private_key = self.__generate_new_private_key()
 
+	# Get the hashed advertisement key (hash of the public key)
     def get_advertisement_key(self) -> str:
         digest = hashes.Hash(hashes.SHA256())
-        digest.update(self
-                      .__derive_elliptic_curve_private_key(self._private_key, ec.SECP224R1())
-                      .public_key()
-                      .public_bytes(Encoding.X962, PublicFormat.CompressedPoint)[1:])
+        digest.update(self.__get_advertisement_key_bytes())
         return base64.b64encode(digest.finalize()).decode()
 
+	# Get the public key
+    def get_public_key(self) -> str:
+        return base64.b64encode(self.__get_advertisement_key_bytes()).decode("ascii")
+
+	# This function works if you follow the standard of the AirTag, as used in OpenHaystack
+    def get_mac_address(self) -> str:
+        adv_key = self.__get_advertisement_key_bytes()
+        first_hex = adv_key[0] | 0b11000000
+        return self.format_byte(first_hex) + ":" + ":".join([self.format_byte(x) for x in adv_key[1:6]])
+
+    @staticmethod
+    def format_byte(byte):
+        return f'{byte:02x}'.upper()
+
+    # Get the X of the public key in bytes
+    def __get_advertisement_key_bytes(self) -> bytes:
+        return self.__derive_elliptic_curve_private_key(self._private_key, ec.SECP224R1()).public_key().public_numbers().x.to_bytes(28, 'big')
+
     @staticmethod
     def __derive_elliptic_curve_private_key(private_key: bytes, curve: ec.EllipticCurve):
         return ec.derive_private_key(int.from_bytes(private_key, 'big'), curve, default_backend())
@@ -52,10 +67,11 @@ def __decrypt_payload(enc_data: bytes, symmetric_key: bytes, tag: bytes):
 
     @staticmethod
     def __decode_tag(data: bytes):
-        latitude = int.from_bytes(data[0:4], 'big', signed=True) / 10000000.0
-        longitude = int.from_bytes(data[4:8], 'big', signed=True) / 10000000.0
-        confidence = int.from_bytes(data[8:9], 'big')
-        return {'lat': latitude, 'lon': longitude, 'conf': confidence}
+        latitude = int.from_bytes(data[0:4], 'big') / 10000000.0
+        longitude = int.from_bytes(data[4:8], 'big') / 10000000.0
+        horizontal_acc = int.from_bytes(data[8:9], 'big')
+        status = data[9]
+        return {'lat': latitude, 'lon': longitude, 'horizontal_acc': horizontal_acc, 'status': status}
 
     @staticmethod
     def __generate_new_private_key():
@@ -64,7 +80,9 @@ def __generate_new_private_key():
 
     def decrypt_message(self, payload):
         data = base64.b64decode(payload)
+        if len(data) > 88: data = data[0:4] + data[5:]
         timestamp = int.from_bytes(data[0:4], 'big')
+        confidence = data[4]
         eph_key = data[5:62]
         shared_key = self.__derive_shared_key_from_private_key_and_eph_key(eph_key)
         derived_key = self.__kdf(shared_key, eph_key)
@@ -74,4 +92,5 @@ def decrypt_message(self, payload):
 
         ret = self.__decode_tag(decrypted)
         ret['timestamp'] = timestamp + 978307200  # 978307200 is delta between unix and cocoa timestamps
+        ret['confidence'] = confidence
         return ret