package services

import (
	"context"
	"fmt"
	"net"
	"strings"
	"time"

	"github.com/shadow1ng/fscan/common"
	"github.com/shadow1ng/fscan/plugins"
)

type TelnetPlugin struct {
	plugins.BasePlugin
}

func NewTelnetPlugin() *TelnetPlugin {
	return &TelnetPlugin{
		BasePlugin: plugins.NewBasePlugin("telnet"),
	}
}

func (p *TelnetPlugin) Scan(ctx context.Context, info *common.HostInfo) *plugins.Result {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)

	if common.DisableBrute {
		return p.identifyService(ctx, info)
	}

	// 构建凭据列表
	credentials := plugins.GenerateCredentials("telnet")
	
	for _, cred := range credentials {
		// 检查上下文是否已取消
		select {
		case <-ctx.Done():
			return &plugins.Result{
				Success: false,
				Service: "telnet",
				Error:   ctx.Err(),
			}
		default:
		}

		if p.testTelnetCredential(ctx, info, cred) {
			common.LogSuccess(fmt.Sprintf("Telnet %s %s:%s", target, cred.Username, cred.Password))
			return &plugins.Result{
				Success:  true,
				Service:  "telnet",
				Username: cred.Username,
				Password: cred.Password,
			}
		}
	}

	return &plugins.Result{
		Success: false,
		Service: "telnet",
		Error:   fmt.Errorf("未发现弱密码"),
	}
}

// testTelnetCredential 测试telnet凭据
func (p *TelnetPlugin) testTelnetCredential(ctx context.Context, info *common.HostInfo, cred plugins.Credential) bool {
	address := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	// 创建带超时的连接
	conn, err := net.DialTimeout("tcp", address, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		return false
	}
	defer conn.Close()
	
	// 设置超时
	deadline := time.Now().Add(time.Duration(common.Timeout) * time.Second)
	conn.SetDeadline(deadline)
	
	// 简单的telnet认证流程
	return p.performSimpleTelnetAuth(conn, cred.Username, cred.Password)
}

// performSimpleTelnetAuth 执行简单的telnet认证
func (p *TelnetPlugin) performSimpleTelnetAuth(conn net.Conn, username, password string) bool {
	buffer := make([]byte, 1024)
	
	// 处理IAC协商并等待真正的登录提示
	loginPromptReceived := false
	attempts := 0
	maxAttempts := 10 // 最多尝试10次读取
	
	for attempts < maxAttempts && !loginPromptReceived {
		attempts++
		
		// 设置较短的读取超时
		conn.SetReadDeadline(time.Now().Add(2 * time.Second))
		n, err := conn.Read(buffer)
		if err != nil {
			common.LogDebug(fmt.Sprintf("第%d次读取失败: %v", attempts, err))
			time.Sleep(200 * time.Millisecond)
			continue
		}
		
		response := string(buffer[:n])
		
		// 处理IAC协商
		p.handleIACNegotiation(conn, buffer[:n])
		
		// 清理响应
		cleaned := p.cleanResponse(response)
		common.LogDebug(fmt.Sprintf("第%d次响应[%s:%s]: %q -> %q", attempts, username, password, response, cleaned))
		
		// 检查是否为shell提示符（无需认证）
		if p.isShellPrompt(cleaned) {
			common.LogDebug(fmt.Sprintf("检测到shell提示符，无需认证"))
			return true
		}
		
		// 检查是否收到登录提示
		if strings.Contains(strings.ToLower(cleaned), "login") ||
		   strings.Contains(strings.ToLower(cleaned), "username") ||
		   strings.Contains(cleaned, ":") { // 简单的提示符检测
			loginPromptReceived = true
			common.LogDebug(fmt.Sprintf("检测到登录提示"))
			break
		}
		
		time.Sleep(200 * time.Millisecond)
	}
	
	if !loginPromptReceived {
		common.LogDebug(fmt.Sprintf("未在%d次尝试中检测到登录提示", maxAttempts))
		return false
	}
	
	// 发送用户名
	common.LogDebug(fmt.Sprintf("发送用户名: %s", username))
	_, err := conn.Write([]byte(username + "\r\n"))
	if err != nil {
		return false
	}
	
	// 等待密码提示
	time.Sleep(500 * time.Millisecond)
	passwordPromptReceived := false
	attempts = 0
	
	for attempts < 5 && !passwordPromptReceived {
		attempts++
		
		conn.SetReadDeadline(time.Now().Add(2 * time.Second))
		n, err := conn.Read(buffer)
		if err != nil {
			common.LogDebug(fmt.Sprintf("读取密码提示第%d次失败: %v", attempts, err))
			time.Sleep(200 * time.Millisecond)
			continue
		}
		
		response := string(buffer[:n])
		cleaned := p.cleanResponse(response)
		common.LogDebug(fmt.Sprintf("密码提示第%d次响应: %q -> %q", attempts, response, cleaned))
		
		if strings.Contains(strings.ToLower(cleaned), "password") ||
		   strings.Contains(cleaned, ":") {
			passwordPromptReceived = true
			common.LogDebug(fmt.Sprintf("检测到密码提示"))
			break
		}
		
		time.Sleep(200 * time.Millisecond)
	}
	
	if !passwordPromptReceived {
		common.LogDebug(fmt.Sprintf("未检测到密码提示"))
		return false
	}
	
	// 发送密码
	common.LogDebug(fmt.Sprintf("发送密码: %s", password))
	_, err = conn.Write([]byte(password + "\r\n"))
	if err != nil {
		return false
	}
	
	// 检查登录结果
	time.Sleep(1000 * time.Millisecond)
	attempts = 0
	
	for attempts < 5 {
		attempts++
		
		conn.SetReadDeadline(time.Now().Add(2 * time.Second))
		n, err := conn.Read(buffer)
		if err != nil {
			common.LogDebug(fmt.Sprintf("读取登录结果第%d次失败: %v", attempts, err))
			time.Sleep(200 * time.Millisecond)
			continue
		}
		
		response := string(buffer[:n])
		cleaned := p.cleanResponse(response)
		common.LogDebug(fmt.Sprintf("登录结果第%d次响应: %q -> %q", attempts, response, cleaned))
		
		// 检查登录成功或失败
		if p.isLoginSuccess(cleaned) {
			common.LogDebug(fmt.Sprintf("登录成功!"))
			return true
		}
		
		if p.isLoginFailed(cleaned) {
			common.LogDebug(fmt.Sprintf("登录失败!"))
			return false
		}
		
		time.Sleep(200 * time.Millisecond)
	}
	
	common.LogDebug(fmt.Sprintf("无法确定登录结果"))
	return false
}

// handleIACNegotiation 处理IAC协商
func (p *TelnetPlugin) handleIACNegotiation(conn net.Conn, data []byte) {
	for i := 0; i < len(data); i++ {
		if data[i] == 255 && i+2 < len(data) { // IAC
			cmd := data[i+1]
			opt := data[i+2]
			
			// 简单响应策略：拒绝所有选项
			switch cmd {
			case 251: // WILL
				// 回应DONT
				conn.Write([]byte{255, 254, opt})
				common.LogDebug(fmt.Sprintf("IAC响应: DONT %d", opt))
			case 253: // DO
				// 回应WONT
				conn.Write([]byte{255, 252, opt})
				common.LogDebug(fmt.Sprintf("IAC响应: WONT %d", opt))
			}
			i += 2
		}
	}
}

// cleanResponse 清理telnet响应中的IAC命令
func (p *TelnetPlugin) cleanResponse(data string) string {
	var result strings.Builder
	for i := 0; i < len(data); i++ {
		b := data[i]
		// 跳过IAC命令序列
		if b == 255 && i+2 < len(data) {
			i += 2
			continue
		}
		// 保留可打印字符
		if (b >= 32 && b <= 126) || b == '\r' || b == '\n' {
			result.WriteByte(b)
		}
	}
	return result.String()
}

// isShellPrompt 检查是否为shell提示符
func (p *TelnetPlugin) isShellPrompt(data string) bool {
	data = strings.ToLower(data)
	return strings.Contains(data, "$") || 
		   strings.Contains(data, "#") || 
		   strings.Contains(data, ">")
}

// isLoginSuccess 检查登录是否成功
func (p *TelnetPlugin) isLoginSuccess(data string) bool {
	data = strings.ToLower(data)
	
	// 检查成功标识
	if strings.Contains(data, "$") || 
	   strings.Contains(data, "#") || 
	   strings.Contains(data, ">") ||
	   strings.Contains(data, "welcome") ||
	   strings.Contains(data, "last login") {
		return true
	}
	
	return false
}

// isLoginFailed 检查登录是否失败
func (p *TelnetPlugin) isLoginFailed(data string) bool {
	data = strings.ToLower(data)
	
	// 检查失败标识
	if strings.Contains(data, "incorrect") ||
	   strings.Contains(data, "failed") ||
	   strings.Contains(data, "denied") ||
	   strings.Contains(data, "invalid") ||
	   strings.Contains(data, "login:") ||  // 重新出现登录提示
	   strings.Contains(data, "username:") {
		return true
	}
	
	return false
}

func (p *TelnetPlugin) identifyService(ctx context.Context, info *common.HostInfo) *plugins.Result {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)

	conn, err := net.DialTimeout("tcp", target, time.Duration(common.Timeout)*time.Second)
	if err != nil {
		return &plugins.Result{
			Success: false,
			Service: "telnet",
			Error:   err,
		}
	}
	defer conn.Close()

	conn.SetDeadline(time.Now().Add(time.Duration(common.Timeout) * time.Second))

	buffer := make([]byte, 1024)
	n, err := conn.Read(buffer)
	if err != nil {
		return &plugins.Result{
			Success: false,
			Service: "telnet",
			Error:   err,
		}
	}

	data := string(buffer[:n])
	// 清理telnet数据
	cleaned := ""
	for i := 0; i < len(data); i++ {
		b := data[i]
		// 跳过telnet控制字符
		if b == 255 && i+2 < len(data) {
			i += 2
			continue
		}
		// 只保留可打印字符和换行符
		if (b >= 32 && b <= 126) || b == '\r' || b == '\n' {
			cleaned += string(b)
		}
	}
	data = cleaned
	
	var banner string
	if strings.Contains(strings.ToLower(data), "login") || strings.Contains(strings.ToLower(data), "username") {
		banner = "Telnet远程终端服务"
	} else if strings.Contains(data, "$") || strings.Contains(data, "#") || strings.Contains(data, ">") {
		banner = "Telnet远程终端服务 (无认证)"
	} else {
		banner = "Telnet远程终端服务"
	}

	common.LogSuccess(fmt.Sprintf("Telnet %s %s", target, banner))

	return &plugins.Result{
		Success: true,
		Service: "telnet",
		Banner:  banner,
	}
}

func init() {
	plugins.RegisterWithPorts("telnet", func() plugins.Plugin {
		return NewTelnetPlugin()
	}, []int{23, 2323})
}