// Command devserver runs the OpenGFW web UI against synthetic data.
//
// The engine itself only builds on Linux (it needs NFQueue), so this little
// program exists to let the frontend be developed and reviewed anywhere:
//
//	go run ./web/devserver
//
// It serves the embedded UI on :8080 with the password "opengfw".
package main

import (
	"context"
	"fmt"
	"log"
	"math/rand"
	"net"
	"os"
	"os/signal"
	"runtime"
	"syscall"
	"time"

	"github.com/apernet/OpenGFW/analyzer"
	"github.com/apernet/OpenGFW/analyzer/tcp"
	"github.com/apernet/OpenGFW/analyzer/udp"
	"github.com/apernet/OpenGFW/modifier"
	modUDP "github.com/apernet/OpenGFW/modifier/udp"
	"github.com/apernet/OpenGFW/ruleset"
	"github.com/apernet/OpenGFW/ruleset/builtins/geo"
	"github.com/apernet/OpenGFW/web"

	"gopkg.in/yaml.v3"
)

const password = "opengfw"

func main() {
	hub := web.NewHub()
	rm := &memoryRules{}
	if err := rm.init(); err != nil {
		log.Fatal(err)
	}

	srv, err := web.NewServer(web.Config{
		Listen: ":8080",
		Secret: password,
		Hub:    hub,
		Rules:  rm,
		Meta: web.MetaInfo{
			Analyzers: []web.AnalyzerInfo{
				{Name: "http", Proto: "tcp"}, {Name: "tls", Proto: "tcp"},
				{Name: "ssh", Proto: "tcp"}, {Name: "socks", Proto: "tcp"},
				{Name: "trojan", Proto: "tcp"}, {Name: "fet", Proto: "tcp"},
				{Name: "dns", Proto: "udp"}, {Name: "quic", Proto: "udp"},
				{Name: "openvpn", Proto: "udp"}, {Name: "wireguard", Proto: "udp"},
			},
			Modifiers: []string{"dns"},
			Actions:   []string{"allow", "block", "drop", "modify"},
			Functions: []string{"geoip", "geosite", "cidr", "lookup"},
		},
		Info: func() web.Info {
			host, _ := os.Hostname()
			return web.Info{
				Version:   "devserver",
				Platform:  runtime.GOOS + "/" + runtime.GOARCH,
				GoVersion: runtime.Version(),
				Hostname:  host,
				RuleFile:  "rules.yaml (in memory)",
				Config: web.ConfigDigest{
					IOQueueSize: 1024, IORST: true, Workers: 4,
					WorkerQueue: 64, UDPMaxStreams: 4096,
					GeoIP: "geoip.dat", GeoSite: "geosite.dat",
				},
			}
		},
		Geo:  geoData,
		Logf: log.Printf,
	})
	if err != nil {
		log.Fatal(err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	go func() {
		ch := make(chan os.Signal, 1)
		signal.Notify(ch, os.Interrupt, syscall.SIGTERM)
		<-ch
		cancel()
	}()

	for i := 0; i < 4; i++ {
		hub.WorkerStarted()
	}
	go generate(ctx, hub)

	log.Printf("web UI on http://127.0.0.1:8080 (password: %s)", password)
	if err := srv.Run(ctx); err != nil {
		log.Fatal(err)
	}
}

var (
	hosts = []string{
		"www.google.com", "github.com", "cdn.jsdelivr.net", "telegram.org",
		"ads.example.net", "tracker.evil.test", "api.openai.com", "www.wikipedia.org",
		"registry.npmjs.org", "malware.bad.test",
	}
	ips = []string{"1.1.1.1", "8.8.8.8", "93.184.216.34", "104.16.132.229", "2606:4700::6810:84e5"}
)

// generate feeds the hub with plausible looking traffic.
func generate(ctx context.Context, hub *web.Hub) {
	rng := rand.New(rand.NewSource(42))
	ticker := time.NewTicker(120 * time.Millisecond)
	defer ticker.Stop()
	var id int64
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			for n := rng.Intn(6); n >= 0; n-- {
				id++
				udp := rng.Intn(3) == 0
				proto := "tcp"
				if udp {
					proto = "udp"
				}
				hub.StreamNew(proto)

				host := hosts[rng.Intn(len(hosts))]
				props := web.Props{}
				if udp {
					props["dns"] = web.PropMap{
						"qr":        false,
						"questions": []map[string]interface{}{{"name": host, "type": 1}},
					}
				} else {
					props["tls"] = web.PropMap{"req": map[string]interface{}{
						"sni": host, "version": 771,
					}}
				}
				hub.PropUpdate(props)

				info := web.StreamInfo{
					ID:      id,
					Proto:   proto,
					SrcIP:   fmt.Sprintf("192.168.1.%d", 2+rng.Intn(60)),
					SrcPort: uint16(20000 + rng.Intn(40000)),
					DstIP:   ips[rng.Intn(len(ips))],
					DstPort: 443,
					Props:   props,
				}

				switch {
				case rng.Intn(10) == 0:
					hub.RuleLog(info, "log-suspicious")
					hub.StreamAction(info, "block")
				case rng.Intn(12) == 0:
					hub.StreamAction(info, "drop")
				case rng.Intn(14) == 0:
					hub.StreamAction(info, "modify")
				case rng.Intn(30) == 0:
					hub.Error(info, "geoip-rule", "lookup timeout")
				default:
					hub.StreamAction(info, "allow")
				}
			}
		}
	}
}

// memoryRules is an in-memory web.RuleManager. Rules are compiled with the real
// ruleset compiler and the real analyzers, so expression errors show up here
// exactly like they would in the engine; only the "apply" step is faked.
type memoryRules struct {
	raw   string
	rules []web.Rule
}

const seedRules = `- name: block-malware
  action: block
  log: true
  expr: 'tls != nil && tls.req != nil && string(tls.req.sni) endsWith ".bad.test"'
- name: block-ads-dns
  action: drop
  expr: 'dns != nil && any(dns.questions, {.name endsWith "ads.example.net"})'
- name: log-ssh
  log: true
  expr: 'ssh != nil'
`

func (m *memoryRules) init() error {
	rules, err := parse(seedRules)
	if err != nil {
		return err
	}
	m.raw, m.rules = seedRules, rules
	return nil
}

func (m *memoryRules) Path() string { return "rules.yaml" }

func (m *memoryRules) Load() (string, []web.Rule, error) { return m.raw, m.rules, nil }

func (m *memoryRules) Validate(raw string) ([]web.Rule, error) { return parse(raw) }

func (m *memoryRules) Marshal(rules []web.Rule) (string, error) {
	bs, err := yaml.Marshal(rules)
	return string(bs), err
}

func (m *memoryRules) Apply(raw string) ([]web.Rule, error) {
	rules, err := parse(raw)
	if err != nil {
		return nil, err
	}
	m.raw, m.rules = raw, rules
	return rules, nil
}

var (
	analyzers = []analyzer.Analyzer{
		&tcp.FETAnalyzer{}, &tcp.HTTPAnalyzer{}, &tcp.SocksAnalyzer{}, &tcp.SSHAnalyzer{},
		&tcp.TLSAnalyzer{}, &tcp.TrojanAnalyzer{}, &udp.DNSAnalyzer{}, &udp.OpenVPNAnalyzer{},
		&udp.QUICAnalyzer{}, &udp.WireGuardAnalyzer{},
	}
	modifiers = []modifier.Modifier{&modUDP.DNSModifier{}}
)

// parse compiles rules the same way the engine does.
func parse(raw string) ([]web.Rule, error) {
	exprRules, err := ruleset.ExprRulesFromYAMLBytes([]byte(raw))
	if err != nil {
		return nil, fmt.Errorf("failed to parse rules: %w", err)
	}
	_, err = ruleset.CompileExprRules(exprRules, analyzers, modifiers, &ruleset.BuiltinConfig{
		Logger:          nopRulesetLogger{},
		GeoSiteFilename: os.Getenv("OPENGFW_GEOSITE"),
		GeoIpFilename:   os.Getenv("OPENGFW_GEOIP"),
		ProtectedDialContext: func(ctx context.Context, network, address string) (net.Conn, error) {
			return (&net.Dialer{}).DialContext(ctx, network, address)
		},
	})
	if err != nil {
		return nil, err
	}
	out := make([]web.Rule, 0, len(exprRules))
	for _, r := range exprRules {
		wr := web.Rule{Name: r.Name, Action: r.Action, Log: r.Log, Expr: r.Expr}
		if r.Modifier.Name != "" {
			wr.Modifier = &web.RuleModifier{Name: r.Modifier.Name, Args: r.Modifier.Args}
		}
		out = append(out, wr)
	}
	return out, nil
}

type nopRulesetLogger struct{}

func (nopRulesetLogger) Log(ruleset.StreamInfo, string)               {}
func (nopRulesetLogger) MatchError(ruleset.StreamInfo, string, error) {}

// geoData lists the geo databases, if they are available in the working
// directory (or wherever OPENGFW_GEOIP / OPENGFW_GEOSITE point).
func geoData() web.GeoData {
	matcher := geo.NewGeoMatcher(os.Getenv("OPENGFW_GEOSITE"), os.Getenv("OPENGFW_GEOIP"))
	var data web.GeoData
	if entries, err := matcher.ListGeoIP(); err != nil {
		data.IPError = err.Error()
	} else {
		for _, e := range entries {
			data.IP = append(data.IP, web.GeoEntry{Code: e.Code, Count: e.CIDRs})
		}
	}
	if entries, err := matcher.ListGeoSite(); err != nil {
		data.SiteError = err.Error()
	} else {
		for _, e := range entries {
			data.Site = append(data.Site, web.GeoEntry{
				Code: e.Code, Count: e.Domains, Attributes: e.Attributes,
			})
		}
	}
	return data
}