package main

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"log"
	"net/http"
	"os/exec"
	"regexp"
	"runtime"
	"strings"
	"text/template"
)

func execCommand(command string) string {
	var cmd *exec.Cmd
	if runtime.GOOS == "windows" {
		cmd = exec.Command("cmd", "/c", command)
	} else {
		cmd = exec.Command("bash", "-c", command)
	}

	var out bytes.Buffer
	var stderr bytes.Buffer
	cmd.Stdout = &out
	cmd.Stderr = &stderr

	err := cmd.Run()
	if err != nil {
		if stderr.Len() > 0 {
			return fmt.Sprintf("命令执行错误: %s", stderr.String())
		}
		return fmt.Sprintf("执行失败: %v", err)
	}
	return out.String()
}

func b64Decode(encoded string) string {
	decodedBytes, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return "error"
	}
	return string(decodedBytes)
}

func aWAF(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api" {
			next.ServeHTTP(w, r)
			return
		}

		query := r.URL.Query().Get("template")
		if query == "" {
			next.ServeHTTP(w, r)
			return
		}

		blacklist := []string{"ls", "whoami", "cat", "uname", "nc", "flag", "etc", "passwd", "\\*", "pwd", "rm", "cp", "mv", "chmod", "chown", "wget", "curl", "bash", "sh", "python", "perl", "ruby", "system", "eval", "less", "more", "find", "grep", "awk", "sed", "tar", "zip", "unzip", "gzip", "gunzip", "bzip2", "bunzip2", "xz", "unxz", "docker", "kubectl", "git", "svn", "f", "l", "g", ",", "\\?", "&&", "\\|", ";", "`", "\"", ">", "<", ":", "=", "\\(", "\\)", "%", "\\\\", "\\^", "\\$", "!", "@", "#", "&"}
		escaped := make([]string, len(blacklist))
		for i, item := range blacklist {
			escaped[i] = "\\b" + item + "\\b"
		}
		wafRegex := regexp.MustCompile(fmt.Sprintf("(?i)%s", strings.Join(escaped, "|")))

		if wafRegex.MatchString(query) {
			// log.Printf("拦截请求: %s", wafRegex.FindAllString(query, -1))
			http.Error(w, query, 200)
			return
		}

		next.ServeHTTP(w, r)
	})
}

func apiHandler(w http.ResponseWriter, r *http.Request) {
	query := r.URL.Query().Get("template")
	if query == "" {
		http.Error(w, "需要template参数", http.StatusBadRequest)
		return
	}

	funcMap := template.FuncMap{
		"exec":      execCommand,
		"B64Decode": b64Decode,
	}

	tmpl, err := template.New("api").Funcs(funcMap).Parse(query)
	if err != nil {
		http.Error(w, query, http.StatusAccepted)
		return
	}

	var buf bytes.Buffer
	if err := tmpl.Execute(&buf, funcMap); err != nil {
		http.Error(w, query, http.StatusAccepted)
		return
	}

	w.Write(buf.Bytes())
}

func rootHandler(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.NotFound(w, r)
		return
	}

	http.ServeFile(w, r, "index.html")
}

func main() {
	mux := http.NewServeMux()
	mux.HandleFunc("/", rootHandler)
	mux.HandleFunc("/api", apiHandler)

	log.Println("服务器启动在 :80")
	log.Fatal(http.ListenAndServe(":80", aWAF(mux)))
}