masyunis/utils/cmd/iplist2rule/main.go

package main

import (
	"flag"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/TecharoHQ/anubis/lib/config"
	"github.com/facebookgo/flagenv"
	"sigs.k8s.io/yaml"
)

type Rule struct {
	Name       string         `yaml:"name" json:"name"`
	Action     config.Rule    `yaml:"action" json:"action"`
	RemoteAddr []string       `json:"remote_addresses,omitempty" yaml:"remote_addresses,omitempty"`
	Weight     *config.Weight `json:"weight,omitempty" yaml:"weight,omitempty"`
}

func init() {
	flag.Usage = func() {
		fmt.Printf(`Usage of %[1]s:

	%[1]s [flags] <blocklist-url> <filename>

Grabs the contents of the blocklist, converts it to an Anubis ruleset, and writes it to filename.

Flags:
`, filepath.Base(os.Args[0]))

		flag.PrintDefaults()
	}
}

var (
	action         = flag.String("action", "DENY", "Anubis action to take (ALLOW / DENY / WEIGH)")
	manualRuleName = flag.String("rule-name", "", "If set, prefer this name over inferring from filename")
	weight         = flag.Int("weight", 0, "If set to any number, add/subtract this many weight points when --action=WEIGH")
)

func main() {
	flagenv.Parse()
	flag.Parse()

	if flag.NArg() != 2 {
		flag.Usage()
		os.Exit(2)
	}

	blocklistURL := flag.Arg(0)
	foutName := flag.Arg(1)
	ruleName := strings.TrimSuffix(foutName, filepath.Ext(foutName))

	if *manualRuleName != "" {
		ruleName = *manualRuleName
	}

	ruleAction := config.Rule(*action)
	if err := ruleAction.Valid(); err != nil {
		log.Fatalf("--action=%q is invalid: %v", *action, err)
	}

	result := &Rule{
		Name:   ruleName,
		Action: ruleAction,
	}

	if *weight != 0 {
		if ruleAction != config.RuleWeigh {
			log.Fatalf("used --weight=%d but --action=%s", *weight, *action)
		}

		result.Weight = &config.Weight{
			Adjust: *weight,
		}
	}

	ips, err := FetchBlocklist(blocklistURL)
	if err != nil {
		log.Fatalf("can't fetch blocklist %s: %v", blocklistURL, err)
	}

	result.RemoteAddr = ips

	fout, err := os.Create(foutName)
	if err != nil {
		log.Fatalf("can't create output file %q: %v", foutName, err)
	}
	defer fout.Close()

	fmt.Fprintf(fout, "# Generated by %s on %s from %s\n\n", filepath.Base(os.Args[0]), time.Now().Format(time.RFC3339), blocklistURL)

	data, err := yaml.Marshal([]*Rule{result})
	if err != nil {
		log.Fatalf("can't marshal yaml")
	}

	fout.Write(data)
}