#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from impacket.ldap import ldapasn1 as ldapasn1_impacket

def searchResEntry_to_dict(results):
    data = {}
    for attr in results['attributes']:
        key = str(attr['type'])
        value = str(attr['vals'][0])
        data[key] = value
    return data

class CMEModule:
    '''
      Retrieves the different Sites and Subnets of an Active Directory

      Authors:
        Podalirius: @podalirius_
    '''

    def options(self, context, module_options):
        """
            showservers    Toggle printing of servers (default: true)
        """

        self.showservers = True

        if module_options and 'SHOWSERVERS' in module_options:
            if module_options['SHOWSERVERS'].lower() == "true" or module_options['SHOWSERVERS'] == "1":
                self.showservers = True
            elif module_options['SHOWSERVERS'].lower() == "false" or module_options['SHOWSERVERS'] == "0":
                self.showservers = False
            else:
                print("Could not parse showservers option.")

    name = 'subnets'
    description = 'Retrieves the different Sites and Subnets of an Active Directory'
    supported_protocols = ['ldap']
    opsec_safe = True
    multiple_hosts = False

    def on_login(self, context, connection):
        dn = ','.join(["DC=%s" % part for part in connection.domain.split('.')])

        context.log.info('Getting the Sites and Subnets from domain')

        list_sites = connection.ldapConnection.search(
            searchBase="CN=Configuration,%s" % dn,
            searchFilter='(objectClass=site)',
            attributes=['distinguishedName', 'name', 'description'],
            sizeLimit=999
        )
        for site in list_sites:
            if isinstance(site, ldapasn1_impacket.SearchResultEntry) is not True:
                continue
            site = searchResEntry_to_dict(site)
            site_dn = site['distinguishedName']
            site_name = site['name']
            site_description = ""
            if "description" in site.keys():
                site_description = site['description']
            # Getting subnets of this site
            list_subnets = connection.ldapConnection.search(
                searchBase="CN=Sites,CN=Configuration,%s" % dn,
                searchFilter='(siteObject=%s)' % site_dn,
                attributes=['distinguishedName', 'name'],
                sizeLimit=999
            )
            if len([subnet for subnet in list_subnets if isinstance(subnet, ldapasn1_impacket.SearchResultEntry)]) == 0:
                context.log.highlight("Site \"%s\"" % site_name)
            else:
                for subnet in list_subnets:
                    if isinstance(subnet, ldapasn1_impacket.SearchResultEntry) is not True:
                        continue
                    subnet = searchResEntry_to_dict(subnet)
                    subnet_dn = subnet['distinguishedName']
                    subnet_name = subnet['name']

                    if self.showservers:
                        # Getting machines in these subnets
                        list_servers = connection.ldapConnection.search(
                            searchBase=site_dn,
                            searchFilter='(objectClass=server)',
                            attributes=['cn'],
                            sizeLimit=999
                        )
                        if len([server for server in list_servers if isinstance(server, ldapasn1_impacket.SearchResultEntry)]) == 0:
                            if len(site_description) != 0:
                                context.log.highlight("Site \"%s\" (Subnet:%s) (description:\"%s\")" % (site_name, subnet_name, site_description))
                            else:
                                context.log.highlight("Site \"%s\" (Subnet:%s)" % (site_name, subnet_name))
                        else:
                            for server in list_servers:
                                if isinstance(server, ldapasn1_impacket.SearchResultEntry) is not True:
                                    continue
                                server = searchResEntry_to_dict(server)['cn']
                                if len(site_description) != 0:
                                    context.log.highlight("Site \"%s\" (Subnet:%s) (description:\"%s\") (Server:%s)" % (site_name, subnet_name, site_description, server))
                                else:
                                    context.log.highlight("Site \"%s\" (Subnet:%s) (Server:%s)" % (site_name, subnet_name, server))
                    else:
                        if len(site_description) != 0:
                            context.log.highlight("Site \"%s\" (Subnet:%s) (description:\"%s\")" % (site_name, subnet_name, site_description))
                        else:
                            context.log.highlight("Site \"%s\" (Subnet:%s)" % (site_name, subnet_name))