package plugin

import (
	"archive/zip"
	"encoding/json"
	"fmt"
	"github.com/Autumn-27/ScopeSentry-Scan/internal/global"
	"github.com/Autumn-27/ScopeSentry-Scan/internal/options"
	"github.com/Autumn-27/ScopeSentry-Scan/internal/types"
	"github.com/Autumn-27/ScopeSentry-Scan/pkg/logger"
	"github.com/Autumn-27/ScopeSentry-Scan/pkg/utils"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"time"
)

func GetName() string {
	return "AI-Infra-Guard"
}

func Install() error {
	toolPath := filepath.Join(global.ExtDir, "AI-Infra-Guard")
	if err := os.MkdirAll(toolPath, os.ModePerm); err != nil {
		logger.SlogError(fmt.Sprintf("Failed to create AI-Infra-Guard folder:", err))
		return err
	}
	osType := runtime.GOOS
	var downloadURL string
	var fileName string
	switch osType {
	case "windows":
		downloadURL = "https://github.com/Autumn-27/AI-Infra-Guard/releases/download/v0.0.5/AI-Infra-Guard_0.0.5_windows_amd64.zip"
		fileName = "ai-infra-guard.exe"
	case "linux":
		downloadURL = "https://github.com/Autumn-27/AI-Infra-Guard/releases/download/v0.0.5/AI-Infra-Guard_0.0.5_linux_amd64.zip"
		fileName = "ai-infra-guard"
	}
	toolExecPath := filepath.Join(toolPath, fileName)
	if _, err := os.Stat(toolExecPath); os.IsNotExist(err) {
		// 创建目标目录、结果目录
		os.MkdirAll(filepath.Join(toolPath, "target"), os.ModePerm)
		os.MkdirAll(filepath.Join(toolPath, "result"), os.ModePerm)
		downloadPath := filepath.Join(global.ExtDir, "AI-Infra-Guard", "AI-Infra-Guard.zip") // 临时下载路径
		success, err := utils.Tools.HttpGetDownloadFile(downloadURL, downloadPath)
		if err != nil || !success {
			logger.SlogErrorLocal(fmt.Sprintf("Failed to download AI-Infra-Guard: %v", err))
			return err
		}

		logger.SlogInfo("Findomain Download successful")

		err = Unzip(downloadPath, toolPath)
		if err != nil {
			fmt.Printf("Failed to extract AI-Infra-Guard: %v", err)
			return nil
		}
		switch osType {
		case "linux":
			err = os.Chmod(toolExecPath, 0755)
			if err != nil {
				fmt.Sprintf("Failed to set permissions: %v", err)
				return nil
			}
		}
		defer utils.Tools.DeleteFile(downloadPath)
		logger.SlogInfo("AI-Infra-Guard installed successfully")
	}
	return nil
}

func Unzip(src string, dest string) error {
	r, err := zip.OpenReader(src)
	if err != nil {
		return err
	}
	defer r.Close()

	for _, f := range r.File {
		fpath := filepath.Join(dest, f.Name)

		if f.FileInfo().IsDir() {
			// Make Folder
			os.MkdirAll(fpath, os.ModePerm)
			continue
		}

		err := os.MkdirAll(filepath.Dir(fpath), os.ModePerm)
		if err != nil {
			return err
		}

		outFile, err := os.OpenFile(fpath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, f.Mode())
		if err != nil {
			return err
		}

		rc, err := f.Open()
		if err != nil {
			return err
		}

		_, err = io.Copy(outFile, rc)

		// Close the file without defer to close before next iteration of loop
		outFile.Close()
		rc.Close()

		if err != nil {
			return err
		}
	}
	return nil
}

func Check() error {
	return nil
}

func Uninstall() error {
	return nil
}

type HttpResult struct {
	URL           string       `json:"url"`            // Target URL
	Title         string       `json:"title"`          // Page title
	ContentLength int          `json:"content-length"` // Response content length
	StatusCode    int          `json:"status-code"`    // HTTP status code
	ResponseTime  string       `json:"response-time"`  // Request response time
	Fingers       []FpResult   `json:"fingerprints"`   // Fingerprint detection results
	Advisories    []VersionVul `json:"advisories"`     // Vulnerability advisory information
	s             string       // Internal string representation
}

type FpResult struct {
	Name    string `json:"name"`
	Version string `json:"version,omitempty"`
	Type    string `json:"type,omitempty"`
}

type VersionVul struct {
	Info       Info     `yaml:"info"`       // Basic vulnerability information
	Rule       string   `yaml:"rule"`       // Rule expression in string format
	References []string `yaml:"references"` // Reference links
}
type Info struct {
	FingerPrintName string `yaml:"name"`                      // Name of the fingerprint
	CVEName         string `yaml:"cve"`                       // CVE identifier
	Summary         string `yaml:"summary"`                   // Brief summary of the vulnerability
	Details         string `yaml:"details"`                   // Detailed description
	CVSS            string `yaml:"cvss"`                      // CVSS score
	Severity        string `yaml:"severity"`                  // Severity level
	SecurityAdvise  string `yaml:"security_advise,omitempty"` // Security advisory
}

func Execute(input interface{}, op options.PluginOption) (interface{}, error) {
	targets := ""
	switch a := input.(type) {
	case []types.AssetOther:
		return nil, nil
	case []types.AssetHttp:
		for _, assetHttp := range a {
			targets += assetHttp.URL + "\n"
		}
	default:
		return nil, nil
	}
	toolPath := filepath.Join(global.ExtDir, "AI-Infra-Guard")
	scanId := utils.Tools.GenerateRandomString(6)
	osType := runtime.GOOS
	var fileName string
	switch osType {
	case "windows":
		fileName = "ai-infra-guard.exe"
	case "linux":
		fileName = "ai-infra-guard"
	}
	targetPath := filepath.Join(toolPath, "target", scanId)
	resultPath := filepath.Join(toolPath, "result", scanId)
	vulDir := filepath.Join(toolPath, "data", "vuln")
	fingerDir := filepath.Join(toolPath, "data", "fingerprints")
	err := utils.Tools.WriteContentFile(targetPath, targets)
	if err != nil {
		op.Log(fmt.Sprintf("write target error: %v", err))
		return nil, err
	}
	start := time.Now()
	defer utils.Tools.DeleteFile(targetPath)
	defer utils.Tools.DeleteFile(resultPath)
	exePath := filepath.Join(toolPath, fileName)
	args := []string{"-file", targetPath, "-o", resultPath, "-json", "-fps", fingerDir, "-vul", vulDir}
	// 使用有超时时间以及上下文管理的命令执行 方便处理异常以及适配暂停任务
	err = utils.Tools.ExecuteCommandWithTimeout(exePath, args, time.Duration(10)*time.Minute, op.Ctx)
	if err != nil {
		op.Log(fmt.Sprintf("ExecuteCommandWithTimeout error: %v", err), "w")
		return nil, err
	}
	resultChan := make(chan string)
	go utils.Tools.ReadFileLineByLine(resultPath, resultChan, op.Ctx)
	for result := range resultChan {
		if result != "" {
			var p HttpResult
			err := json.Unmarshal([]byte(result), &p)
			if err != nil {
				op.Log(fmt.Sprintf("result to json error: %v %v", err, result), "w")
				continue
			}
			for _, ad := range p.Advisories {
				tmp := types.VulnResult{
					Url:     p.URL,
					VulName: ad.Info.FingerPrintName + "-" + ad.Info.CVEName + "-" + ad.Info.Summary,
					Level:   strings.ToLower(ad.Info.Severity),
					Time:    utils.Tools.GetTimeNow(),
					Request: result,
					Status:  1,
					Tags:    []string{"AI-Infra-Guard"},
				}
				op.Log(fmt.Sprintf("found vul: %v %v", tmp.Url, tmp.VulName))
				op.ResultFunc(tmp)
			}

		}
	}
	end := time.Now()
	duration := end.Sub(start)
	op.Log(fmt.Sprintf("AI-Infra-Guard scan completed, time: %v", duration))
	return nil, nil
}
