#!/usr/bin/env python
####################
#
# Copyright (c) 2022 Dirk-jan Mollema (@_dirkjan)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Get attribute sets from AD schema
#
####################
import sys
import argparse
from uuid import UUID
import pprint

import random
import string
import getpass
from ldap3 import NTLM, Server, Connection, ALL
import ldap3
from ldap3.protocol.microsoft import security_descriptor_control

def print_m(string):
    sys.stderr.write('\033[94m[-]\033[0m %s\n' % (string))

def print_o(string):
    sys.stderr.write('\033[92m[+]\033[0m %s\n' % (string))

def print_f(string):
    sys.stderr.write('\033[91m[!]\033[0m %s\n' % (string))

def main():
    parser = argparse.ArgumentParser(description='Get attribute sets from AD schema')
    parser._optionals.title = "Main options"
    parser._positionals.title = "Required options"

    #Main parameters
    #maingroup = parser.add_argument_group("Main options")
    parser.add_argument("host", metavar='HOSTNAME', help="Hostname/ip or ldap://host:port connection string to connect to")
    parser.add_argument("-u", "--user", metavar='USERNAME', help="DOMAIN\\username for authentication")
    parser.add_argument("-p", "--password", metavar='PASSWORD', help="Password or LM:NTLM hash, will prompt if not specified")

    args = parser.parse_args()

    #Prompt for password if not set
    authentication = None
    if args.user is not None:
        authentication = NTLM
        if not '\\' in args.user:
            print_f('Username must include a domain, use: DOMAIN\\username')
            sys.exit(1)
        if args.password is None:
            args.password = getpass.getpass()

    controls = security_descriptor_control(sdflags=0x04)
    # define the server and the connection
    s = Server(args.host, get_info=ALL)
    print_m('Connecting to host...')
    c = Connection(s, user=args.user, password=args.password, authentication=authentication)
    print_m('Binding to host')
    # perform the Bind operation
    if not c.bind():
        print_f('Could not bind with specified credentials')
        print_f(c.result)
        sys.exit(1)
    print_o('Bind OK')

    sresult = c.extend.standard.paged_search('CN=Extended-Rights,'+ c.server.info.other['configurationNamingContext'][0],
                                                     '(rightsGuid=*)',
                                                     attributes=['name', 'rightsGuid', 'displayName'])
    objecttype_guid_map = {}
    guid_objecttype_map = {}
    sets = {}
    for res in sresult:
        if res['attributes']['rightsGuid']:
            guid = res['attributes']['rightsGuid'].lower()
            attname = res['attributes']['displayName']
            name = res['attributes']['name']
            guid_objecttype_map[guid] = (name, attname)

    sresult = c.extend.standard.paged_search(c.server.info.other['schemaNamingContext'][0],
                                                     '(objectClass=*)',
                                                     attributes=['name', 'schemaidguid', 'attributeSecurityGUID'])
    # pprint.pprint(guid_objecttype_map)
    # return
    for res in sresult:
        if res['attributes']['schemaIDGUID']:
            guid = str(UUID(bytes_le=res['attributes']['schemaIDGUID']))
            attname = res['attributes']['name'].lower()
            objecttype_guid_map[attname] = guid
            setguid_bin = res['attributes']['attributeSecurityGUID']
            if setguid_bin:
                setguid = str(UUID(bytes_le=setguid_bin))
            try:
                sets[setguid].append(attname)
            except KeyError:
                sets[setguid] = [attname]
    # pprint.pprint(objecttype_guid_map)
    # pprint.pprint(sets)
    for setguid, setcontent in sets.items():
        try:
            print("{0} [{1}] (GUID: {2}".format(guid_objecttype_map[setguid][0], guid_objecttype_map[setguid][1], setguid ))
        except KeyError:
            print(f'Unknown property set {setguid}')
        pprint.pprint(setcontent)
if __name__ == '__main__':
    main()