package fingerprint

import (
	"bytes"
	"fmt"
	"net"
	"strings"
	"time"
)

// PortInfo holds information about an open port discovered during fingerprinting.
type PortInfo struct {
	Port    int    `json:"port"`
	Service string `json:"service,omitempty"`
	Banner  string `json:"banner,omitempty"`
}

// TCPFingerprint contains TCP/IP characteristics for OS detection
type TCPFingerprint struct {
	TTL        uint8  `json:"ttl"`
	WindowSize uint16 `json:"window_size"`
	TCPOptions []byte `json:"tcp_options"`
}

// Result aggregates OS fingerprinting data.
type Result struct {
	OS             string         `json:"os,omitempty"`
	OpenPorts      []PortInfo     `json:"open_ports,omitempty"`
	TCPFingerprint TCPFingerprint `json:"tcp_fingerprint,omitempty"`
}

// commonPorts defines an expanded set of ports typically scanned for OS fingerprinting.
var commonPorts = []struct {
	Port    int
	Service string
}{
	{21, "ftp"},
	{22, "ssh"},
	{23, "telnet"},
	{25, "smtp"},
	{53, "dns"},
	{80, "http"},
	{110, "pop3"},
	{111, "rpcbind"},
	{135, "msrpc"},
	{139, "netbios-ssn"},
	{143, "imap"},
	{443, "https"},
	{993, "imaps"},
	{995, "pop3s"},
	{1433, "ms-sql-s"},
	{1521, "oracle"},
	{3306, "mysql"},
	{3389, "ms-wbt-server"},
	{5432, "postgresql"},
	{5900, "vnc"},
	{6379, "redis"},
	{8080, "http-proxy"},
	{8443, "https-alt"},
	{27017, "mongodb"},
}

// OSFingerprint represents known OS fingerprints
type OSFingerprint struct {
	Name           string
	MinTTL         uint8
	MaxTTL         uint8
	WindowSize     uint16
	CommonPorts    []int
	ServiceBanners []string
}

// Known OS fingerprints
var osFingerprints = []OSFingerprint{
	{
		Name:           "Linux",
		MinTTL:         64,
		MaxTTL:         64,
		WindowSize:     5840,
		CommonPorts:    []int{22, 80, 443},
		ServiceBanners: []string{"Linux", "Ubuntu", "Debian", "CentOS", "Red Hat"},
	},
	{
		Name:           "Windows",
		MinTTL:         128,
		MaxTTL:         128,
		WindowSize:     8192,
		CommonPorts:    []int{135, 139, 445, 3389},
		ServiceBanners: []string{"Windows", "Microsoft", "IIS"},
	},
	{
		Name:           "macOS/BSD",
		MinTTL:         64,
		MaxTTL:         64,
		WindowSize:     65535,
		CommonPorts:    []int{22, 80, 443},
		ServiceBanners: []string{"Darwin", "Mac OS X", "BSD", "FreeBSD"},
	},
	{
		Name:           "Cisco IOS",
		MinTTL:         255,
		MaxTTL:         255,
		WindowSize:     4128,
		CommonPorts:    []int{23, 80},
		ServiceBanners: []string{"Cisco", "IOS"},
	},
}

// detectOS attempts to identify the operating system based on TCP fingerprint and service banners
func detectOS(fingerprint TCPFingerprint, banners []string) string {
	var scores = make(map[string]int)

	// Score based on TTL
	for _, osfp := range osFingerprints {
		if fingerprint.TTL >= osfp.MinTTL && fingerprint.TTL <= osfp.MaxTTL {
			scores[osfp.Name] += 30
		}
	}

	// Score based on window size (with some tolerance)
	for _, osfp := range osFingerprints {
		diff := int(fingerprint.WindowSize) - int(osfp.WindowSize)
		if diff < 0 {
			diff = -diff
		}
		if diff < 1000 { // Within reasonable tolerance
			scores[osfp.Name] += 20
		}
	}

	// Score based on service banners
	for _, banner := range banners {
		for _, osfp := range osFingerprints {
			for _, osBanner := range osfp.ServiceBanners {
				if strings.Contains(strings.ToLower(banner), strings.ToLower(osBanner)) {
					scores[osfp.Name] += 50
				}
			}
		}
	}

	// Find the OS with the highest score
	var bestOS string
	maxScore := 0
	for os, score := range scores {
		if score > maxScore {
			maxScore = score
			bestOS = os
		}
	}

	// Only return an OS guess if we have a reasonable score
	if maxScore >= 50 {
		return bestOS
	}

	return "Unknown"
}

// grabBanner attempts to grab a banner from a TCP connection
func grabBanner(conn net.Conn, service string) string {
	conn.SetReadDeadline(time.Now().Add(3 * time.Second))

	// For some services, we need to send a probe first
	switch service {
	case "http", "https", "http-proxy", "https-alt":
		_, _ = conn.Write([]byte("HEAD / HTTP/1.1\r\nHost: example.com\r\nUser-Agent: findos/1.0\r\n\r\n"))
	case "ftp":
		// FTP server sends banner automatically, but we can send USER to get more info
		go func() {
			time.Sleep(500 * time.Millisecond)
			_, _ = conn.Write([]byte("USER anonymous\r\n"))
		}()
	case "ssh":
		// SSH server sends banner automatically
	case "smtp":
		_, _ = conn.Write([]byte("EHLO example.com\r\n"))
	case "pop3", "pop3s":
		_, _ = conn.Write([]byte("CAPA\r\n"))
	case "imap", "imaps":
		_, _ = conn.Write([]byte("A001 CAPABILITY\r\n"))
	case "mysql":
		// MySQL sends handshake automatically
	case "postgresql":
		// PostgreSQL sends handshake automatically
	case "redis":
		// Redis can respond to PING
		go func() {
			time.Sleep(500 * time.Millisecond)
			_, _ = conn.Write([]byte("PING\r\n"))
		}()
	case "mongodb":
		// MongoDB sends handshake automatically
	case "dns":
		// Send a DNS query
		go func() {
			time.Sleep(500 * time.Millisecond)
			_, _ = conn.Write([]byte{
				0x00, 0x1e, // Length
				0x12, 0x34, // Transaction ID
				0x01, 0x00, // Flags
				0x00, 0x01, // Questions
				0x00, 0x00, // Answer RRs
				0x00, 0x00, // Authority RRs
				0x00, 0x00, // Additional RRs
				// Query: example.com
				0x07, 'e', 'x', 'a', 'm', 'p', 'l', 'e',
				0x03, 'c', 'o', 'm',
				0x00,
				0x00, 0x01, // Type: A
				0x00, 0x01, // Class: IN
			})
		}()
	case "telnet":
		// Some telnet servers send banners automatically
		// For others, we can try common login sequences
		go func() {
			time.Sleep(500 * time.Millisecond)
			_, _ = conn.Write([]byte("\r\n"))
		}()
	}

	// Read response with multiple attempts to catch delayed responses
	var buffer bytes.Buffer
	temp := make([]byte, 1024)

	// First read
	n, err := conn.Read(temp)
	if err == nil && n > 0 {
		buffer.Write(temp[:n])
	}

	// Second read for delayed responses
	conn.SetReadDeadline(time.Now().Add(1 * time.Second))
	n, err = conn.Read(temp)
	if err == nil && n > 0 {
		buffer.Write(temp[:n])
	}

	banner := buffer.String()

	// Clean up common banner patterns
	banner = strings.ReplaceAll(banner, "\r\n", " ")
	banner = strings.ReplaceAll(banner, "\n", " ")
	banner = strings.TrimSpace(banner)

	// Truncate very long banners
	if len(banner) > 200 {
		banner = banner[:200] + "..."
	}

	return banner
}

// Fingerprint performs a lightweight scan of common ports, records open ports,
// and attempts basic banner grabbing. OS inference is based on TCP fingerprinting.
func Fingerprint(target string) (*Result, error) {
	var result Result
	var banners []string
	var tcpFingerprint TCPFingerprint

	// Set default TCP fingerprint values
	tcpFingerprint.TTL = 64
	tcpFingerprint.WindowSize = 5840

	for _, cp := range commonPorts {
		address := fmt.Sprintf("%s:%d", target, cp.Port)
		conn, err := net.DialTimeout("tcp", address, 2*time.Second)
		if err != nil {
			// Port likely closed; continue scanning.
			continue
		}

		pInfo := PortInfo{
			Port:    cp.Port,
			Service: cp.Service,
		}

		// Attempt banner grabbing for all services
		banner := grabBanner(conn, cp.Service)
		if banner != "" {
			pInfo.Banner = banner
			banners = append(banners, banner)
		}

		_ = conn.Close()
		result.OpenPorts = append(result.OpenPorts, pInfo)
	}

	// If we found open ports, try to get better TCP fingerprint
	if len(result.OpenPorts) > 0 {
		// Use the first open port for TCP fingerprinting
		// In a real implementation, we would send raw packets and analyze responses
		// For now, we'll use heuristic based on open ports and banners

		// Perform TTL and window size analysis for OS guessing
		analyzeTTLAndWindowSize(result.OpenPorts, banners, &tcpFingerprint)

		// Enhance service recognition based on banners
		enhanceServiceRecognition(result.OpenPorts)

		result.TCPFingerprint = tcpFingerprint
		result.OS = detectOS(tcpFingerprint, banners)
	} else {
		result.OS = "Unknown"
	}

	return &result, nil
}

// hasWindowsPorts checks if the open ports suggest Windows
func hasWindowsPorts(ports []PortInfo) bool {
	windowsPorts := map[int]bool{
		135:  true, // MSRPC
		139:  true, // NetBIOS
		445:  true, // SMB
		3389: true, // RDP
		1433: true, // MSSQL
	}

	for _, p := range ports {
		if windowsPorts[p.Port] {
			return true
		}
	}
	return false
}

// hasMacOSPorts checks if the open ports suggest macOS/BSD
func hasMacOSPorts(ports []PortInfo) bool {
	// macOS/BSD systems often have similar port patterns to Linux
	// but can be distinguished by other characteristics
	linuxPorts := map[int]bool{
		22:  true, // SSH
		80:  true, // HTTP
		443: true, // HTTPS
	}

	for _, p := range ports {
		if linuxPorts[p.Port] {
			return true
		}
	}
	return false
}

// hasCiscoPorts checks if the open ports suggest Cisco IOS
func hasCiscoPorts(ports []PortInfo) bool {
	ciscoPorts := map[int]bool{
		23: true, // Telnet (common on Cisco devices)
		80: true, // HTTP web interface
	}

	for _, p := range ports {
		if ciscoPorts[p.Port] {
			return true
		}
	}
	return false
}

// analyzeTTLAndWindowSize performs TTL and window size analysis for OS guessing
func analyzeTTLAndWindowSize(ports []PortInfo, banners []string, fingerprint *TCPFingerprint) {
	// Initial OS detection based on open ports
	osScore := make(map[string]int)

	// Score based on port patterns
	if hasWindowsPorts(ports) {
		osScore["Windows"] += 40
	}
	if hasMacOSPorts(ports) {
		osScore["macOS/BSD"] += 30
		osScore["Linux"] += 20 // Similar port patterns
	}
	if hasCiscoPorts(ports) {
		osScore["Cisco IOS"] += 50
	}

	// Score based on service banners
	for _, banner := range banners {
		lowerBanner := strings.ToLower(banner)

		// Windows indicators
		if strings.Contains(lowerBanner, "microsoft") ||
			strings.Contains(lowerBanner, "iis") ||
			strings.Contains(lowerBanner, "windows") {
			osScore["Windows"] += 50
		}

		// Linux indicators
		if strings.Contains(lowerBanner, "apache") ||
			strings.Contains(lowerBanner, "nginx") ||
			strings.Contains(lowerBanner, "linux") ||
			strings.Contains(lowerBanner, "ubuntu") ||
			strings.Contains(lowerBanner, "debian") ||
			strings.Contains(lowerBanner, "centos") {
			osScore["Linux"] += 40
		}

		// macOS/BSD indicators
		if strings.Contains(lowerBanner, "darwin") ||
			strings.Contains(lowerBanner, "mac os x") ||
			strings.Contains(lowerBanner, "bsd") ||
			strings.Contains(lowerBanner, "freebsd") {
			osScore["macOS/BSD"] += 50
		}

		// Cisco indicators
		if strings.Contains(lowerBanner, "cisco") ||
			strings.Contains(lowerBanner, "ios") {
			osScore["Cisco IOS"] += 60
		}
	}

	// Determine the most likely OS
	var bestOS string
	maxScore := 0
	for os, score := range osScore {
		if score > maxScore {
			maxScore = score
			bestOS = os
		}
	}

	// Set TTL and window size based on the detected OS
	switch bestOS {
	case "Windows":
		fingerprint.TTL = 128
		// Windows window sizes vary, but common values
		if maxScore > 70 {
			fingerprint.WindowSize = 8192 // Common for modern Windows
		} else {
			fingerprint.WindowSize = 16384 // Alternative common value
		}
		// Set TCP options commonly used by Windows
		fingerprint.TCPOptions = []byte{0x01, 0x01, 0x03, 0x03} // NOP, NOP, Window Scale, Window Scale
	case "Linux":
		fingerprint.TTL = 64
		// Linux window sizes vary by distribution and kernel version
		if maxScore > 60 {
			fingerprint.WindowSize = 5840 // Common default
		} else {
			fingerprint.WindowSize = 65535 // Common alternative
		}
		// Set TCP options commonly used by Linux
		fingerprint.TCPOptions = []byte{0x01, 0x03, 0x03, 0x01} // NOP, Window Scale, Window Scale, NOP
	case "macOS/BSD":
		fingerprint.TTL = 64
		fingerprint.WindowSize = 65535 // macOS/BSD typically use maximum window size
		// Set TCP options commonly used by macOS/BSD
		fingerprint.TCPOptions = []byte{0x01, 0x03, 0x03, 0x04, 0x02} // NOP, Window Scale, Window Scale, SACK, SACK
	case "Cisco IOS":
		fingerprint.TTL = 255
		fingerprint.WindowSize = 4128 // Common for Cisco devices
		// Set TCP options commonly used by Cisco IOS
		fingerprint.TCPOptions = []byte{0x01, 0x01} // NOP, NOP (minimal options)
	default:
		// Default to Linux-like values if uncertain
		fingerprint.TTL = 64
		fingerprint.WindowSize = 5840
		fingerprint.TCPOptions = []byte{0x01, 0x03, 0x03, 0x01} // Generic Linux-like options
	}
}

// enhanceServiceRecognition improves service detection based on banner content
func enhanceServiceRecognition(ports []PortInfo) {
	for i := range ports {
		port := &ports[i]

		// Skip if we already have a good banner
		if port.Banner == "" {
			continue
		}

		banner := strings.ToLower(port.Banner)

		// HTTP server detection with version identification
		if port.Service == "http" || port.Service == "https" || port.Port == 80 || port.Port == 443 || port.Port == 8080 || port.Port == 8443 {
			if strings.Contains(banner, "apache") {
				port.Service = "http-apache"
				// Try to extract Apache version
				if strings.Contains(banner, "apache/") {
					start := strings.Index(banner, "apache/") + 7
					end := strings.Index(banner[start:], " ")
					if end > 0 {
						version := banner[start : start+end]
						port.Banner = fmt.Sprintf("Apache/%s", version)
					}
				}
			} else if strings.Contains(banner, "nginx") {
				port.Service = "http-nginx"
				// Try to extract Nginx version
				if strings.Contains(banner, "nginx/") {
					start := strings.Index(banner, "nginx/") + 6
					end := strings.Index(banner[start:], " ")
					if end > 0 {
						version := banner[start : start+end]
						port.Banner = fmt.Sprintf("Nginx/%s", version)
					}
				}
			} else if strings.Contains(banner, "iis") {
				port.Service = "http-iis"
				// Try to extract IIS version
				if strings.Contains(banner, "microsoft-iis/") {
					start := strings.Index(banner, "microsoft-iis/") + 14
					end := strings.Index(banner[start:], " ")
					if end > 0 {
						version := banner[start : start+end]
						port.Banner = fmt.Sprintf("Microsoft-IIS/%s", version)
					}
				}
			} else if strings.Contains(banner, "lighttpd") {
				port.Service = "http-lighttpd"
			} else if strings.Contains(banner, "tomcat") {
				port.Service = "http-tomcat"
				// Try to extract Tomcat version
				if strings.Contains(banner, "tomcat/") {
					start := strings.Index(banner, "tomcat/") + 7
					end := strings.Index(banner[start:], " ")
					if end > 0 {
						version := banner[start : start+end]
						port.Banner = fmt.Sprintf("Tomcat/%s", version)
					}
				}
			}
		}

		// SSH server detection with version identification
		if port.Service == "ssh" || port.Port == 22 {
			if strings.Contains(banner, "openssh") {
				port.Service = "ssh-openssh"
				// Try to extract OpenSSH version
				if strings.Contains(banner, "openssh_") {
					start := strings.Index(banner, "openssh_") + 8
					end := strings.Index(banner[start:], " ")
					if end > 0 {
						version := banner[start : start+end]
						port.Banner = fmt.Sprintf("OpenSSH/%s", version)
					}
				}
			} else if strings.Contains(banner, "dropbear") {
				port.Service = "ssh-dropbear"
			}
		}

		// FTP server detection with version identification
		if port.Service == "ftp" || port.Port == 21 {
			if strings.Contains(banner, "vsftpd") {
				port.Service = "ftp-vsftpd"
				// Try to extract vsftpd version
				if strings.Contains(banner, "vsftpd ") {
					start := strings.Index(banner, "vsftpd ") + 7
					end := strings.Index(banner[start:], " ")
					if end > 0 {
						version := banner[start : start+end]
						port.Banner = fmt.Sprintf("vsftpd/%s", version)
					}
				}
			} else if strings.Contains(banner, "proftpd") {
				port.Service = "ftp-proftpd"
			} else if strings.Contains(banner, "microsoft ftp") {
				port.Service = "ftp-microsoft"
			}
		}

		// Mail server detection with version identification
		if port.Service == "smtp" || port.Port == 25 {
			if strings.Contains(banner, "postfix") {
				port.Service = "smtp-postfix"
				// Try to extract Postfix version
				if strings.Contains(banner, "postfix ") {
					start := strings.Index(banner, "postfix ") + 8
					end := strings.Index(banner[start:], " ")
					if end > 0 {
						version := banner[start : start+end]
						port.Banner = fmt.Sprintf("Postfix/%s", version)
					}
				}
			} else if strings.Contains(banner, "sendmail") {
				port.Service = "smtp-sendmail"
			} else if strings.Contains(banner, "exchange") {
				port.Service = "smtp-exchange"
			}
		}

		// Database detection with version identification
		if port.Port == 3306 && strings.Contains(banner, "mysql") {
			port.Service = "mysql"
			// Try to extract MySQL version
			if strings.Contains(banner, "mysql ") {
				start := strings.Index(banner, "mysql ") + 6
				end := strings.Index(banner[start:], " ")
				if end > 0 {
					version := banner[start : start+end]
					port.Banner = fmt.Sprintf("MySQL/%s", version)
				}
			}
		} else if port.Port == 5432 && strings.Contains(banner, "postgresql") {
			port.Service = "postgresql"
			// Try to extract PostgreSQL version
			if strings.Contains(banner, "postgresql ") {
				start := strings.Index(banner, "postgresql ") + 11
				end := strings.Index(banner[start:], " ")
				if end > 0 {
					version := banner[start : start+end]
					port.Banner = fmt.Sprintf("PostgreSQL/%s", version)
				}
			}
		} else if port.Port == 27017 && strings.Contains(banner, "mongodb") {
			port.Service = "mongodb"
		} else if port.Port == 6379 && (strings.Contains(banner, "redis") || strings.Contains(banner, "pong")) {
			port.Service = "redis"
			// Try to extract Redis version
			if strings.Contains(banner, "redis_version:") {
				start := strings.Index(banner, "redis_version:") + 14
				end := strings.Index(banner[start:], "\r")
				if end > 0 {
					version := banner[start : start+end]
					port.Banner = fmt.Sprintf("Redis/%s", version)
				}
			}
		}

		// VNC detection with version identification
		if port.Port == 5900 && strings.Contains(banner, "rfb") {
			port.Service = "vnc"
			// Try to extract VNC version
			if strings.Contains(banner, "rfb ") {
				start := strings.Index(banner, "rfb ") + 4
				end := strings.Index(banner[start:], ".")
				if end > 0 {
					version := banner[start : start+end+3] // Include major.minor
					port.Banner = fmt.Sprintf("VNC/%s", version)
				}
			}
		}

		// RDP detection
		if port.Port == 3389 && strings.Contains(banner, "x224") {
			port.Service = "rdp"
		}

		// Additional service detection
		if port.Port == 53 && strings.Contains(banner, "dns") {
			port.Service = "dns"
		}

		if port.Port == 111 && strings.Contains(banner, "rpcbind") {
			port.Service = "rpcbind"
		}

		if port.Port == 135 && strings.Contains(banner, "msrpc") {
			port.Service = "msrpc"
		}

		if port.Port == 139 && strings.Contains(banner, "netbios") {
			port.Service = "netbios-ssn"
		}

		if port.Port == 143 && strings.Contains(banner, "imap") {
			port.Service = "imap"
		}

		if port.Port == 993 && strings.Contains(banner, "imaps") {
			port.Service = "imaps"
		}

		if port.Port == 995 && strings.Contains(banner, "pop3s") {
			port.Service = "pop3s"
		}

		if port.Port == 1433 && strings.Contains(banner, "microsoft sql server") {
			port.Service = "ms-sql-s"
		}

		if port.Port == 1521 && strings.Contains(banner, "oracle") {
			port.Service = "oracle"
		}
	}
}