cleanup-orgs

#!/usr/bin/env python3

import argparse

import mongoengine as me
from mist.api.users.methods import purge_org
from mist.api.models import User, Organization


def main():
    argparser = argparse.ArgumentParser(
        description='Cleanup broken references from orgs'
    )
    argparser.add_argument('-o', '--organization',
                           help='ID of org to cleanup.',
                           required=False)
    args = argparser.parse_args()

    def cleanup_org(org):
        for member in org.members:
            if member.__class__ != User:
                if len(org.members) > 1:
                    org.members.remove(member)
                    try:
                        org.save()
                    except me.ValidationError:
                        print('Purging org %s' % org.name)
                        purge_org(org)
                        print('Done purging org')
                else:
                    print('Purging org %s' % org.name)
                    purge_org(org)
                    print('Done purging org')

    if not args.organization:
        for org in Organization.objects:
            cleanup_org(org)
    else:
        org = Organization.objects.get(id=args.organization)
        cleanup_org(org)


if __name__ == '__main__':
    main()