fix script entry point and #28

bloodhound_import/__init__.py

@@ -1,4 +1,5 @@
 import argparse
+import asyncio
 import logging
 from bloodhound_import import database
 from bloodhound_import.importer import parse_file, add_constraints
@@ -51,4 +52,4 @@ async def main():
 
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())

bloodhound_import/database.py

@@ -2,13 +2,13 @@
 Code borrowed from
 https://github.com/fox-it/aclpwn.py/blob/master/aclpwn/database.py
 """
-from neo4j import AsyncGraphDatabase
+from neo4j import AsyncGraphDatabase, AsyncDriver
 import platform
 import json
 import os
 
 
-def init_driver(ip, port, scheme, user, password):
+def init_driver(ip, port, scheme, user, password) -> AsyncDriver:
     uri = "{}://{}:{}".format(scheme, ip, port)
     driver = AsyncGraphDatabase.driver(uri, auth=(user, password))
     return driver

bloodhound_import/importer.py

@@ -343,7 +343,7 @@ async def parse_container(tx: neo4j.Transaction, container: dict):
             query = build_add_edge_query('Container', target['ObjectType'], 'Contains', '{isacl: false}')
             await tx.run(query, props=dict(source=identifier, target=target['ObjectIdentifier']))
 
-async def parse_file(filename: str, driver: neo4j.GraphDatabase):
+async def parse_file(filename: str, driver: neo4j.AsyncDriver):
     """Parse a bloodhound file.
 
     Arguments:

setup.py

@@ -16,6 +16,6 @@
     packages=['bloodhound_import'],
     install_requires=['neo4j-driver', 'ijson'],
     entry_points={
-        'console_scripts': ['bloodhound-import=bloodhound_import:main']
+        'console_scripts': ['bloodhound-import=bloodhound_import:__main__']
     }
 )