Add grabbing key from remote AEA

fixes #8

get_key.py

@@ -25,41 +25,40 @@ def error(msg):
     exit(1)
 
 
-def main(aea_path: Path, verbose: bool = False):
+def get_key(f, verbose: bool = False):
     fields = {}
-    with aea_path.open("rb") as f:
-        header = f.read(12)
-        if len(header) != 12:
-            error(f"Expected 12 bytes, got {len(header)}")
+    header = f.read(12)
+    if len(header) != 12:
+        error(f"Expected 12 bytes, got {len(header)}")
 
-        magic = header[:4]
-        if magic != b"AEA1":
-            error(f"Invalid magic: {magic.hex()}")
+    magic = header[:4]
+    if magic != b"AEA1":
+        error(f"Invalid magic: {magic.hex()}")
 
-        profile = int.from_bytes(header[4:7], "little")
-        if profile != AEA_PROFILE__HKDF_SHA256_AESCTR_HMAC__SYMMETRIC__NONE:
-            error(f"Invalid AEA profile: {profile}")
+    profile = int.from_bytes(header[4:7], "little")
+    if profile != AEA_PROFILE__HKDF_SHA256_AESCTR_HMAC__SYMMETRIC__NONE:
+        error(f"Invalid AEA profile: {profile}")
 
-        auth_data_blob_size = int.from_bytes(header[8:12], "little")
+    auth_data_blob_size = int.from_bytes(header[8:12], "little")
 
-        if auth_data_blob_size == 0:
-            error("No auth data blob")
+    if auth_data_blob_size == 0:
+        error("No auth data blob")
 
-        auth_data_blob = f.read(auth_data_blob_size)
-        if len(auth_data_blob) != auth_data_blob_size:
-            error(f"Expected {auth_data_blob_size} bytes, got {len(auth_data_blob)}")
+    auth_data_blob = f.read(auth_data_blob_size)
+    if len(auth_data_blob) != auth_data_blob_size:
+        error(f"Expected {auth_data_blob_size} bytes, got {len(auth_data_blob)}")
 
-        assert auth_data_blob[:4]
+    assert auth_data_blob[:4]
 
-        while len(auth_data_blob) > 0:
-            field_size = int.from_bytes(auth_data_blob[:4], "little")
-            field_blob = auth_data_blob[:field_size]
+    while len(auth_data_blob) > 0:
+        field_size = int.from_bytes(auth_data_blob[:4], "little")
+        field_blob = auth_data_blob[:field_size]
 
-            key, value = field_blob[4:].split(b"\x00", 1)
+        key, value = field_blob[4:].split(b"\x00", 1)
 
-            fields[key.decode()] = value.decode()
+        fields[key.decode()] = value.decode()
 
-            auth_data_blob = auth_data_blob[field_size:]
+        auth_data_blob = auth_data_blob[field_size:]
 
     if verbose:
         pprint(fields, stream=sys.stderr)
@@ -89,10 +88,26 @@ def main(aea_path: Path, verbose: bool = False):
         print(base64.b64encode(pt).decode())
 
 
+def main(path: str, verbose: bool = False):
+    if path.startswith("http://") or path.startswith("https://"):
+        with requests.get(path, timeout=10, stream=True) as response:
+            response.raise_for_status()
+
+            get_key(response.raw, verbose)
+
+    else:
+        aea_path = Path(path)
+        if not aea_path.exists():
+            error(f"File {path} does not exist")
+
+        with aea_path.open("rb") as f:
+            get_key(f, verbose)
+
+
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Get the key for an AEA file")
-    parser.add_argument("path", help="Path to the AEA file")
+    parser = argparse.ArgumentParser(description="Get the key for an AEA file or URL")
+    parser.add_argument("path", help="Path or URL to the AEA file")
     parser.add_argument("-v", "--verbose", action="store_true", help="Show verbose output")
     args = parser.parse_args()
 
-    main(Path(args.path), args.verbose)
+    main(args.path, args.verbose)