package checker

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"time"

	sdk "git.happydns.org/checker-sdk-go/checker"
)

const malwareBazaarEndpoint = "https://mb-api.abuse.ch/api/v1/"

func init() { Register(&malwareBazaarSource{endpoint: malwareBazaarEndpoint}) }

type malwareBazaarSource struct {
	endpoint string
}

func (*malwareBazaarSource) ID() string   { return "malwarebazaar" }
func (*malwareBazaarSource) Name() string { return "abuse.ch MalwareBazaar" }

func (*malwareBazaarSource) Options() SourceOptions {
	return SourceOptions{
		Admin: []sdk.CheckerOptionField{
			{
				Id:          "malwarebazaar_auth_key",
				Type:        "string",
				Label:       "MalwareBazaar Auth-Key",
				Description: "abuse.ch MalwareBazaar Auth-Key (free, requires an abuse.ch account). Without this key the source is disabled.",
				Secret:      true,
			},
		},
		User: []sdk.CheckerOptionField{
			{
				Id:          "enable_malwarebazaar",
				Type:        "bool",
				Label:       "Use abuse.ch MalwareBazaar",
				Description: "Search MalwareBazaar for malware samples tagged with the domain (typically C2 infrastructure or delivery hosts).",
				Default:     true,
			},
		},
	}
}

func (s *malwareBazaarSource) Query(ctx context.Context, domain, registered string, opts sdk.CheckerOptions) []SourceResult {
	authKey := stringOpt(opts, "malwarebazaar_auth_key")
	if !sdk.GetBoolOption(opts, "enable_malwarebazaar", true) || registered == "" || authKey == "" {
		return []SourceResult{{SourceID: s.ID(), SourceName: s.Name(), Enabled: false}}
	}

	res := SourceResult{
		SourceID: s.ID(), SourceName: s.Name(), Enabled: true,
		Reference: "https://bazaar.abuse.ch/browse/",
	}

	buf, err := json.Marshal(map[string]any{
		"query": "search_tag",
		"tag":   registered,
	})
	if err != nil {
		res.Error = err.Error()
		return []SourceResult{res}
	}

	reqCtx, cancel := context.WithTimeout(ctx, 15*time.Second)
	defer cancel()

	req, err := http.NewRequestWithContext(reqCtx, http.MethodPost, s.endpoint, bytes.NewReader(buf))
	if err != nil {
		res.Error = err.Error()
		return []SourceResult{res}
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Auth-Key", authKey)

	body, status, err := httpDo(req, 4<<20)
	if err != nil {
		res.Error = err.Error()
		return []SourceResult{res}
	}
	if status != http.StatusOK {
		res.Error = fmt.Sprintf("HTTP %d: %s", status, truncate(string(body), 200))
		return []SourceResult{res}
	}

	var parsed struct {
		QueryStatus string `json:"query_status"`
		Data        []struct {
			SHA256Hash string `json:"sha256_hash"`
			FileName   string `json:"file_name"`
			FileType   string `json:"file_type_mime"`
			Signature  string `json:"signature"`
			FirstSeen  string `json:"first_seen"`
		} `json:"data"`
	}
	if err := json.Unmarshal(body, &parsed); err != nil {
		res.Error = "decode: " + err.Error()
		return []SourceResult{res}
	}

	switch parsed.QueryStatus {
	case "ok":
		signatures := map[string]bool{}
		for _, sample := range parsed.Data {
			if sample.Signature != "" && !signatures[sample.Signature] {
				signatures[sample.Signature] = true
				res.Reasons = append(res.Reasons, sample.Signature)
			}
			res.Evidence = append(res.Evidence, Evidence{
				Label:  "Sample",
				Value:  sample.SHA256Hash,
				Status: sample.FileType,
				Extra: map[string]string{
					"filename":  sample.FileName,
					"signature": sample.Signature,
					"first_seen": sample.FirstSeen,
				},
			})
		}
	case "no_results", "illegal_search_term":
		// Clean.
	default:
		res.Error = "query_status=" + parsed.QueryStatus
	}
	return []SourceResult{res}
}

func (*malwareBazaarSource) Evaluate(r SourceResult) (bool, string) {
	if r.Enabled && r.Error == "" && len(r.Evidence) > 0 {
		return true, SeverityWarn
	}
	return false, ""
}

func (*malwareBazaarSource) Diagnose(res SourceResult) Diagnosis {
	return Diagnosis{
		Severity: SeverityWarn,
		Title:    "Associated with malware samples in abuse.ch MalwareBazaar",
		Detail: fmt.Sprintf(
			"%d malware sample(s) are tagged with this domain; malware family/signature(s): %s. MalwareBazaar samples are tagged with their C2 domain or delivery host when known. Investigate the domain's hosting and DNS records to identify and remove malicious infrastructure.",
			len(res.Evidence), joinNonEmpty(res.Reasons, ", "),
		),
		Fix:      res.Reference,
		FixIsURL: res.Reference != "",
	}
}