Workaround for SpecterOps/BloodHound-Legacy#614

sharphound-import-fix.py

#!/usr/bin/env python3
from argparse import ArgumentParser, RawDescriptionHelpFormatter
from zipfile import ZipFile
import json
import shutil

# from https://github.com/fox-it/BloodHound.py/blob/273984883d9ca3dd21f6fca35ec88671cac3fc87/bloodhound/ad/trusts.py#L59
TRUST_DIRECTIONS = {
    0: 'Disabled',
    1: 'Inbound',
    2: 'Outbound',
    3: 'Bidirectional',
}
TRUST_TYPES = {
    0: 'ParentChild',
    1: 'CrossLink',
    2: 'Forest',
    3: 'External',
    4: 'Unknown',
}
ENCODING = 'utf-8-sig'


def main() -> None:
    entrypoint = ArgumentParser(
        formatter_class=RawDescriptionHelpFormatter,
        epilog=(
            'examples:\n'
            '  python3 ./sharphound-import-fix.py ./*_BloodHound.zip ./*_domains.json\n'
        ),
    )
    entrypoint.add_argument('file', nargs='+', metavar='ZIPFILE|JSONFILE...')
    opts = entrypoint.parse_args()

    for path in opts.file:
        if path.endswith('.json'):
            update_json(path)
            print(f'fixed {path}')
        elif path.endswith('.zip'):
            update_zip(path)
            print(f'fixed {path}')
        else:
            print(f'warning: unknown file extension: {path}')


def update_json(path: str) -> None:
    with open(path, 'r', encoding=ENCODING) as file:
        data = json.load(file)
    data = fix(data)
    with open(path, 'w', encoding=ENCODING) as file:
        file.write(json.dumps(data, indent=2, sort_keys=False))


def update_zip(path: str) -> None:
    with ZipFile(path, 'r') as zipin:
        with ZipFile(f'{path}.tmp', 'w') as zipout:
            for entry in zipin.infolist():
                if entry.filename.endswith('_domains.json'):
                    data = json.loads(zipin.read(entry).decode(ENCODING))
                    data = fix(data)
                    zipout.writestr(entry, json.dumps(data).encode(ENCODING))
                else:
                    zipout.writestr(entry, zipin.read(entry))
    shutil.move(f'{path}.tmp', path)


def fix(data: dict) -> dict:
    for item in data['data']:
        for trust in item['Trusts']:
            if trust['TrustDirection'] not in TRUST_DIRECTIONS.values():
                trust['TrustDirection'] = TRUST_DIRECTIONS[trust['TrustDirection']]
            if trust['TrustType'] not in TRUST_TYPES.values():
                trust['TrustType'] = TRUST_TYPES[trust['TrustType']]
    return data


if __name__ == '__main__':
    main()