package query

import (
	"Yasso/pkg/exploit/ldap/core/filter"
	"encoding/binary"
	"fmt"
	"github.com/audibleblink/bamflags"
	"github.com/bwmarrin/go-objectsid"
	"github.com/jedib0t/go-pretty/v6/table"
	"gopkg.in/ldap.v2"
	"math"
	"strconv"
	"strings"
)

// LdapListResolver ldap查询解析器
func LdapListResolver(t table.Writer, conn *ldap.Conn, baseDN string, command string) error {
	res, err := conn.SearchWithPaging(ldap.NewSearchRequest(
		baseDN,
		ldap.ScopeWholeSubtree,
		ldap.NeverDerefAliases,
		math.MaxInt32,
		0,
		false,
		filter.LdapFilter("*", LdapQueries[command]),
		[]string{},
		nil,
	), math.MaxInt32)
	if err != nil {
		return err
	}
	// 查询到数据
	if len(res.Entries) > 0 {
		for _, m := range res.Entries {
			t.AppendRows([]table.Row{
				{fmt.Sprintf(" - %s", m.GetAttributeValue("sAMAccountName"))},
			})
		}
		t.AppendSeparator()
		if command == "users" {
			t.AppendRows([]table.Row{
				{strconv.Itoa(len(res.Entries)) + "Domain Users Found"},
			})
		}
		if command == "computers" {
			t.AppendRows([]table.Row{
				{strconv.Itoa(len(res.Entries)) + " Domain Computers Found"},
			})
		}
		if command == "groups" {
			t.AppendRows([]table.Row{
				{strconv.Itoa(len(res.Entries)) + " Domain Groups Found"},
			})
		}
	}
	return nil
}

// LdapFullResolver ldap的完全解析器
func LdapFullResolver(t table.Writer, conn *ldap.Conn, baseDN string, command string) error {
	res, err := conn.SearchWithPaging(ldap.NewSearchRequest(
		baseDN,
		ldap.ScopeWholeSubtree,
		ldap.NeverDerefAliases,
		math.MaxInt32,
		0,
		false,
		filter.LdapFilter("*", LdapQueries[command]),
		[]string{},
		nil,
	), math.MaxInt32)
	if err != nil {
		return nil
	}
	// 查询到数据
	if len(res.Entries) > 0 {
		for _, m := range res.Entries {
			LdapResolver(m, t)
			t.AppendSeparator()
		}
	}
	return nil
}

// LdapSpecificFullResolver ldap的特殊解析器
func LdapSpecificFullResolver(t table.Writer, conn *ldap.Conn, baseDN string, name string, command string) error {
	res, err := conn.SearchWithPaging(ldap.NewSearchRequest(
		baseDN,
		ldap.ScopeWholeSubtree,
		ldap.NeverDerefAliases,
		math.MaxInt32,
		0,
		false,
		filter.LdapFilter("*", strings.ReplaceAll(LdapQueries[fmt.Sprintf("specific-%s", command)], "{SAM}", name)),
		[]string{},
		nil,
	), math.MaxInt32)
	if err != nil {
		return err
	}
	if len(res.Entries) > 0 {
		for _, m := range res.Entries {
			LdapResolver(m, t)
			t.AppendSeparator()
		}
	}
	return nil
}

func LdapUserMemberShipResolver(t table.Writer, conn *ldap.Conn, baseDN string, name string, command string) error {
	res, err := conn.SearchWithPaging(ldap.NewSearchRequest(
		baseDN,
		ldap.ScopeWholeSubtree,
		ldap.NeverDerefAliases,
		math.MaxInt32,
		0,
		false,
		filter.LdapFilter("*", strings.ReplaceAll(LdapQueries[fmt.Sprintf("specific-%s", command)], "{SAM}", name)),
		[]string{},
		nil,
	), math.MaxInt32)
	if err != nil {
		return err
	}
	if len(res.Entries) > 0 {
		for _, m := range res.Entries {
			if len(m.GetAttributeValue("memberOf")) > 0 {
				t.AppendRows([]table.Row{
					{fmt.Sprintf("Member of: \n\t%s", strings.Join(m.GetAttributeValues("memberOf"), "\n\t"))},
				})
			}
		}
	}
	return nil
}

func LdapGroupDN(conn *ldap.Conn, baseDN string, name string) string {
	res, err := conn.SearchWithPaging(ldap.NewSearchRequest(
		baseDN,
		ldap.ScopeWholeSubtree,
		ldap.NeverDerefAliases,
		math.MaxInt32,
		0,
		false,
		filter.LdapFilter("*", strings.ReplaceAll(LdapQueries["specific-groups"], "{SAM}", name)),
		[]string{},
		nil,
	), math.MaxInt32)
	if err != nil {
		fmt.Println(err)
	}
	if len(res.Entries) > 0 {
		for _, m := range res.Entries {
			return m.DN
		}
	}
	return ""
}

func LdapGroupResolver(t table.Writer, conn *ldap.Conn, baseDN string, name string) error {
	result, err := conn.SearchWithPaging(ldap.NewSearchRequest(
		baseDN,
		ldap.ScopeWholeSubtree,
		ldap.NeverDerefAliases,
		math.MaxInt32,
		0,
		false,
		filter.LdapFilter("*", strings.ReplaceAll(LdapQueries["group-members"], "{DN}", LdapGroupDN(conn, baseDN, name))),
		[]string{},
		nil,
	), math.MaxInt32)
	if err != nil {
		fmt.Println(err)
	}
	if len(result.Entries) > 0 {
		t.AppendRows([]table.Row{
			{"Domain users in group: "},
		})
		for _, m := range result.Entries {
			t.AppendRows([]table.Row{
				{fmt.Sprintf(" - %s", m.GetAttributeValue("sAMAccountName"))},
			})

		}
	}
	return nil
}
func LdapParse(attr int64, Map map[int]string) []string {
	values, err := bamflags.ParseInt(attr)
	if err != nil {
		return []string{}
	}
	var all []string
	for _, value := range values {
		if propName := Map[value]; propName != "" {
			all = append(all, propName)
		}
	}
	return all
}

func LdapGUID(b []byte) (string, error) {
	if len(b) != 16 {
		return "", fmt.Errorf("guid must be 16 bytes")
	}
	return fmt.Sprintf(
		"%08x-%04x-%04x-%04x-%012x",
		binary.LittleEndian.Uint32(b[:4]),
		binary.LittleEndian.Uint16(b[4:6]),
		binary.LittleEndian.Uint16(b[6:8]),
		b[8:10],
		b[10:]), nil
}

func LdapSID(b []byte) (string, error) {
	if len(b) < 12 {
		return "", fmt.Errorf("[+]Invalid Windows SID")
	}
	sid := objectsid.Decode(b)
	return sid.String(), nil
}

// LdapResolver ldap的解析器
func LdapResolver(entry *ldap.Entry, t table.Writer) {
	// DN
	if len(entry.DN) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("DN: %s", entry.DN)},
		})
	}
	// 用户名
	if len(entry.GetAttributeValue("sAMAccountName")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("SAM Account Name: %s", entry.GetAttributeValue("sAMAccountName"))},
		})
	}
	// 用户名类型
	if len(entry.GetAttributeValue("sAMAccountType")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("SAM Account Type: %s", entry.GetAttributeValue("sAMAccountType"))},
		})
	}
	// Mail
	if len(entry.GetAttributeValue("mail")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Mail ID: %s", entry.GetAttributeValue("mail"))},
		})
	}
	// UID
	if len(entry.GetAttributeValue("uid")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("UID: %s", entry.GetAttributeValue("uid"))},
		})
	}
	// OU
	if len(entry.GetAttributeValue("ou")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("OU: %s", entry.GetAttributeValue("ou"))},
		})
	}
	// CN
	if len(entry.GetAttributeValue("cn")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("CN: %s", entry.GetAttributeValue("cn"))},
		})
	}
	if len(entry.GetAttributeValue("givenName")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Given Name: %s", entry.GetAttributeValue("givenName"))},
		})
	}
	if len(entry.GetAttributeValue("sn")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("SN: %s", entry.GetAttributeValue("sn"))},
		})
	}
	if len(entry.GetAttributeValue("description")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Description: %s", entry.GetAttributeValue("description"))},
		})
	}
	if len(entry.GetAttributeValue("title")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Title: %s", entry.GetAttributeValue("title"))},
		})
	}
	if len(entry.GetAttributeValue("c")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Country: %s", entry.GetAttributeValue("c"))},
		})
	}
	if len(entry.GetAttributeValue("co")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Country Code: %s", entry.GetAttributeValue("co"))},
		})
	}
	if len(entry.GetAttributeValue("l")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("City: %s", entry.GetAttributeValue("l"))},
		})
	}
	if len(entry.GetAttributeValue("st")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("State: %s", entry.GetAttributeValue("st"))},
		})
	}
	if len(entry.GetAttributeValue("streetAddress")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Street Address: %s", entry.GetAttributeValue("streetAddress"))},
		})
	}
	if len(entry.GetAttributeValue("postalCode")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Postal Code: %s", entry.GetAttributeValue("postalCode"))},
		})
	}
	if len(entry.GetAttributeValue("postOfficeBox")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Post Office Box: %s", entry.GetAttributeValue("postOfficeBox"))},
		})
	}
	if len(entry.GetAttributeValue("company")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Company: %s", entry.GetAttributeValue("company"))},
		})
	}
	if len(entry.GetAttributeValue("instanceType")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Instance Type: %s", entry.GetAttributeValue("instanceType"))},
		})
	}
	if len(entry.GetAttributeValue("objectClass")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Object Class: %s", strings.Join(entry.GetAttributeValues("objectClass"), ", "))},
		})
	}
	if len(entry.GetAttributeValue("objectCategory")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Object Category: %s", entry.GetAttributeValue("objectCategory"))},
		})
	}
	if len(entry.GetAttributeValue("memberOf")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Member Of: \n\t%s", strings.Join(entry.GetAttributeValues("memberOf"), "\n\t"))},
		})
	}
	if len(entry.GetAttributeValue("dNSHostName")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("DNS Host Name: %s", entry.GetAttributeValue("dNSHostName"))},
		})
	}
	if len(entry.GetAttributeValue("servicePrincipalName")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Service Principal Name (SPN): \n\t%s", strings.Join(entry.GetAttributeValues("servicePrincipalName"), "\n\t"))},
		})
	}
	if len(entry.GetAttributeValue("operatingSystem")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Operating System: %s", entry.GetAttributeValue("operatingSystem"))},
		})
	}
	if len(entry.GetAttributeValue("operatingSystemVersion")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("Operating System Version: %s", entry.GetAttributeValue("operatingSystemVersion"))},
		})
	}
	if len(entry.GetAttributeValue("userAccountControl")) > 0 {
		uacFlag, err := strconv.ParseInt(entry.GetAttributeValue("userAccountControl"), 0, 64)
		if err != nil {
			//do nothing
		} else {
			t.AppendRows([]table.Row{
				{fmt.Sprintf("UAC Flag: %s", strings.Join(LdapParse(uacFlag, UacFlags)[:], ","))},
			})
		}
	}
	if len(entry.GetAttributeValue("userAccountControl")) > 0 {
		t.AppendRows([]table.Row{
			{fmt.Sprintf("User Account Control: %s", entry.GetAttributeValue("userAccountControl"))},
		})
	}
	if len(entry.GetAttributeValue("trustType")) > 0 {
		trustType, err := strconv.ParseInt(entry.GetAttributeValue("trustType"), 0, 64)
		if err != nil {
			//do nothing
		} else {
			t.AppendRows([]table.Row{
				{fmt.Sprintf("Trust Type: %s", strings.Join(LdapParse(trustType, TrustType)[:], ","))},
			})
		}
	}
	if len(entry.GetAttributeValue("trustDirection")) > 0 {
		trustDirection, err := strconv.ParseInt(entry.GetAttributeValue("trustDirection"), 0, 64)
		if err != nil {
			//do nothing
		} else {
			t.AppendRows([]table.Row{
				{fmt.Sprintf("Trust Direction: %s", strings.Join(LdapParse(trustDirection, TrustDirections)[:], ","))},
			})
		}
	}
	if len(entry.GetAttributeValue("trustAttributes")) > 0 {
		trustAttributes, err := strconv.ParseInt(entry.GetAttributeValue("trustAttributes"), 0, 64)
		if err != nil {
			//do nothing
		} else {
			t.AppendRows([]table.Row{
				{fmt.Sprintf("Trust Attribute: %s", strings.Join(LdapParse(trustAttributes, Trust)[:], ","))},
			})
		}
	}
	if len(entry.GetAttributeValue("objectGUID")) > 0 {
		guid, err := LdapGUID(entry.GetRawAttributeValue("objectGUID"))
		if err != nil {
			//do nothing
		} else {
			t.AppendRows([]table.Row{
				{fmt.Sprintf("Object GUID: %s", guid)},
			})
		}
	}
	if len(entry.GetAttributeValue("objectSid")) > 0 {
		sidByte := []byte(entry.GetAttributeValue("objectSid"))
		sid, err := LdapSID(sidByte)
		if err != nil {
			//do nothing
		} else {
			t.AppendRows([]table.Row{
				{fmt.Sprintf("Object SID: %s", sid)},
			})
		}
	}
}