package modbus

import (
	"context"
	"encoding/binary"
	"fmt"
	"net"
	"time"

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

// ModbusPlugin Modbus插件实现
type ModbusPlugin struct {
	*base.ServicePlugin
	exploiter *ModbusExploiter
}

// NewModbusPlugin 创建Modbus插件
func NewModbusPlugin() *ModbusPlugin {
	// 插件元数据
	metadata := &base.PluginMetadata{
		Name:        "modbus",
		Version:     "2.0.0",
		Author:      "fscan-team",
		Description: "Modbus工业协议扫描和利用插件",
		Category:    "service",
		Ports:       []int{502, 5020}, // 502: 标准Modbus TCP端口, 5020: 测试端口
		Protocols:   []string{"tcp"},
		Tags:        []string{"modbus", "industrial", "scada", "unauthorized"},
	}
	
	// 创建连接器和服务插件
	connector := NewModbusConnector()
	servicePlugin := base.NewServicePlugin(metadata, connector)
	
	// 创建Modbus插件
	plugin := &ModbusPlugin{
		ServicePlugin: servicePlugin,
		exploiter:     NewModbusExploiter(),
	}
	
	// 设置能力
	plugin.SetCapabilities([]base.Capability{
		base.CapUnauthorized,
		base.CapDataExtraction,
	})
	
	return plugin
}

// Scan 重写扫描方法，检测Modbus协议和未授权访问
func (p *ModbusPlugin) Scan(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	// 如果禁用了暴力破解，只进行服务识别
	if common.DisableBrute {
		return p.performServiceIdentification(ctx, info)
	}
	
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	// Modbus协议通常无认证，直接检查协议响应
	unauthCred := &base.Credential{Username: "", Password: ""}
	result, err := p.ScanCredential(ctx, info, unauthCred)
	if err == nil && result.Success {
		// 获取设备信息
		deviceInfo := p.getDeviceInfo(ctx, info)
		
		// 未授权访问成功
		common.LogSuccess(i18n.GetText("modbus_unauth_access", target))
		if deviceInfo != "" {
			common.LogSuccess(i18n.GetText("modbus_device_info", deviceInfo))
		}
		
		return &base.ScanResult{
			Success:     true,
			Service:     "Modbus",
			Credentials: []*base.Credential{unauthCred},
			Extra: map[string]interface{}{
				"service":      "Modbus",
				"port":         info.Ports,
				"unauthorized": true,
				"access_type":  "no_authentication",
				"device_info":  deviceInfo,
			},
		}, nil
	}
	
	// 如果Modbus协议检测失败，返回失败结果
	return &base.ScanResult{
		Success: false,
		Error:   fmt.Errorf("非Modbus服务或连接失败"),
	}, nil
}

// generateCredentials Modbus通常无需凭据，返回空凭据
func (p *ModbusPlugin) generateCredentials() []*base.Credential {
	// Modbus协议通常无认证机制，只返回空凭据用于协议检测
	return []*base.Credential{
		{Username: "", Password: ""},
	}
}

// getDeviceInfo 获取Modbus设备信息
func (p *ModbusPlugin) getDeviceInfo(ctx context.Context, info *common.HostInfo) string {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	timeout := time.Duration(common.Timeout) * time.Second
	
	// 尝试建立连接
	conn, err := net.DialTimeout("tcp", target, timeout)
	if err != nil {
		return ""
	}
	defer conn.Close()
	
	// 设置超时
	conn.SetDeadline(time.Now().Add(timeout))
	
	// 发送读取设备标识请求 (功能码0x11)
	request := buildModbusDeviceIdentifyRequest()
	if _, err := conn.Write(request); err != nil {
		return ""
	}
	
	// 读取响应
	response := make([]byte, 256)
	n, err := conn.Read(response)
	if err != nil {
		return ""
	}
	
	return parseDeviceInfo(response[:n])
}

// buildModbusDeviceIdentifyRequest 构建Modbus设备标识请求
func buildModbusDeviceIdentifyRequest() []byte {
	request := make([]byte, 8)
	
	// Modbus TCP头部
	binary.BigEndian.PutUint16(request[0:], 0x0002) // 事务标识符
	binary.BigEndian.PutUint16(request[2:], 0x0000) // 协议标识符
	binary.BigEndian.PutUint16(request[4:], 0x0002) // 长度
	request[6] = 0x01                               // 单元标识符
	request[7] = 0x11                               // 功能码: Report Server ID
	
	return request
}

// parseDeviceInfo 解析设备信息
func parseDeviceInfo(response []byte) string {
	if len(response) < 8 {
		return "Unknown Device"
	}
	
	// 检查是否为有效响应
	if !isValidModbusResponse(response) {
		return "Unknown Device"
	}
	
	unitID := response[6]
	funcCode := response[7]
	
	info := fmt.Sprintf("Unit ID: %d, Function: 0x%02X", unitID, funcCode)
	
	// 如果是设备标识响应，尝试解析设备信息
	if funcCode == 0x11 && len(response) > 9 {
		byteCount := response[8]
		if byteCount > 0 && len(response) >= 9+int(byteCount) {
			// 提取设备ID信息
			deviceData := response[9 : 9+int(byteCount)]
			if len(deviceData) > 0 {
				info += fmt.Sprintf(", Device Data: %x", deviceData)
			}
		}
	} else if funcCode == 0x01 && len(response) >= 10 {
		// 读取线圈响应，显示线圈状态
		byteCount := response[8]
		if byteCount > 0 && len(response) >= 9+int(byteCount) {
			coilValue := response[9] & 0x01
			info += fmt.Sprintf(", Coil Status: %d", coilValue)
		}
	}
	
	return info
}

// Exploit 使用exploiter执行利用
func (p *ModbusPlugin) Exploit(ctx context.Context, info *common.HostInfo, creds *base.Credential) (*base.ExploitResult, error) {
	return p.exploiter.Exploit(ctx, info, creds)
}

// GetExploitMethods 获取利用方法
func (p *ModbusPlugin) GetExploitMethods() []base.ExploitMethod {
	return p.exploiter.GetExploitMethods()
}

// IsExploitSupported 检查利用支持
func (p *ModbusPlugin) IsExploitSupported(method base.ExploitType) bool {
	return p.exploiter.IsExploitSupported(method)
}

// performServiceIdentification 执行Modbus服务识别（-nobr模式）
func (p *ModbusPlugin) performServiceIdentification(ctx context.Context, info *common.HostInfo) (*base.ScanResult, error) {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	
	// 尝试识别Modbus服务
	modbusInfo, isModbus := p.identifyModbusService(ctx, info)
	if isModbus {
		// 记录服务识别成功
		common.LogSuccess(i18n.GetText("modbus_service_identified", target, modbusInfo))
		
		return &base.ScanResult{
			Success: true,
			Service: "Modbus",
			Banner:  modbusInfo,
			Extra: map[string]interface{}{
				"service": "Modbus",
				"port":    info.Ports,
				"info":    modbusInfo,
			},
		}, nil
	}
	
	// 如果无法识别为Modbus，返回失败
	return &base.ScanResult{
		Success: false,
		Error:   fmt.Errorf("无法识别为Modbus服务"),
	}, nil
}

// identifyModbusService 通过协议识别Modbus服务
func (p *ModbusPlugin) identifyModbusService(ctx context.Context, info *common.HostInfo) (string, bool) {
	target := fmt.Sprintf("%s:%s", info.Host, info.Ports)
	timeout := time.Duration(common.Timeout) * time.Second
	
	// 尝试建立TCP连接
	conn, err := net.DialTimeout("tcp", target, timeout)
	if err != nil {
		return "", false
	}
	defer conn.Close()
	
	// 设置操作超时
	conn.SetDeadline(time.Now().Add(timeout))
	
	// 发送Modbus读取线圈请求
	request := buildModbusReadCoilsRequest()
	if _, err := conn.Write(request); err != nil {
		return "", false
	}
	
	// 读取响应
	response := make([]byte, 256)
	n, err := conn.Read(response)
	if err != nil {
		return "", false
	}
	
	// 检查是否为有效Modbus响应
	if isValidModbusResponse(response[:n]) {
		deviceInfo := parseDeviceInfo(response[:n])
		return fmt.Sprintf("Modbus TCP服务: %s", deviceInfo), true
	}
	
	return "", false
}

// =============================================================================
// 插件注册
// =============================================================================

// RegisterModbusPlugin 注册Modbus插件
func RegisterModbusPlugin() {
	factory := base.NewSimplePluginFactory(
		&base.PluginMetadata{
			Name:        "modbus",
			Version:     "2.0.0",
			Author:      "fscan-team",
			Description: "Modbus工业协议扫描和利用插件",
			Category:    "service",
			Ports:       []int{502, 5020}, // 502: 标准Modbus TCP端口, 5020: 测试端口
			Protocols:   []string{"tcp"},
			Tags:        []string{"modbus", "industrial", "scada", "unauthorized"},
		},
		func() base.Plugin {
			return NewModbusPlugin()
		},
	)
	
	base.GlobalPluginRegistry.Register("modbus", factory)
}

// 自动注册
func init() {
	RegisterModbusPlugin()
}