长亭雷池 WAF 优化拦截响应内容

package main

import (
    "github.com/fsnotify/fsnotify"
    "log"
    "os"
    "strings"
    "time"
)

// SAFELINE_DIR 监听的目录
var SAFELINE_DIR = os.Getenv("SAFELINE_DIR")

// Content 需要强制修改的文件内容
var Content = map[string]string{
    "/data/forbidden_pages/default_forbidden_page.html": `{"code": 0, "msg": "风控拦截"}`,
}

func init() {
    if strings.HasSuffix(SAFELINE_DIR, "/") {
        SAFELINE_DIR = SAFELINE_DIR[:len(SAFELINE_DIR)-1]
    }

    // 初始化阶段强制修改一次
    go func() {
        time.Sleep(time.Second * 30)

        for path := range Content {
            _ = forceContent(path)
        }
    }()
}

func main() {
    err := handle()

    if err != nil {
        log.Fatal(err)
    }
}

func handle() error {
    watcher, err := fsnotify.NewWatcher()

    if err != nil {
        return err
    }

    defer watcher.Close()

    go func() {
        for {
            select {
            case event, ok := <-watcher.Events:
                if !ok {
                    return
                }

                // 只监听写入事件
                if !event.Has(fsnotify.Write) {
                    continue
                }

                err = forceContent(event.Name)

                if err != nil {
                    watcher.Errors <- err
                }
            case err, ok := <-watcher.Errors:
                if !ok {
                    return
                }

                log.Println("error:", err)
            }
        }
    }()

    err = watcher.Add(SAFELINE_DIR)

    if err != nil {
        return err
    }

    <-make(chan struct{})

    return nil
}

func forceContent(path string) error {
    log.Printf("强制修改文件: %s", path)

    content, ok := Content[path]

    if !ok {
        return nil
    }

    body, err := os.ReadFile(path)

    if err != nil {
        return err
    }

    if string(body) == content {
        return nil
    }

    return os.WriteFile(path, []byte(content), 0644)
}