counters/cmd/counterctl/main.go

// counterctl shows today's (or a given day's) incoming/outgoing counters for
// an address or a domain by asking the running migadu_counters daemon
// (GET /counters), the same way rcptctl talks to rcptd.
package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/BurntSushi/toml"

	"migadu_counters/internal/api"
	"migadu_counters/internal/config"
)

// Set by -ldflags at build time.
var (
	version = "dev"
	commit  = "none"
	date    = "unknown"
)

func main() {
	os.Exit(run(os.Args[1:], os.Stdout, os.Stderr, os.Getenv, time.Now))
}

// cliConfig mirrors the daemon config fields the CLI needs; TOML ignores the
// rest, so the daemon's own config file can be used as is.
type cliConfig struct {
	Addr      string `toml:"addr"`
	AuthToken string `toml:"auth_token"`
}

func run(args []string, stdout, stderr io.Writer, getenv func(string) string, now func() time.Time) int {
	fs := flag.NewFlagSet("counterctl", flag.ContinueOnError)
	fs.SetOutput(stderr)
	address := fs.String("address", "", "mailbox address to show (A, D and O levels)")
	domain := fs.String("domain", "", "domain to show (D and O levels)")
	dateFlag := fs.String("date", "", "ISO 8601 calendar day YYYY-MM-DD, UTC (default: today)")
	jsonOut := fs.Bool("json", false, "print the raw API response")
	inOnly := fs.Bool("in", false, "show only the incoming counters")
	outOnly := fs.Bool("out", false, "show only the outgoing counters")
	configPath := fs.String("config", "", "config file (default: "+strings.Join(config.DefaultPaths, " or ")+")")
	urlFlag := fs.String("url", "", "service base URL (default: derived from addr in the config)")
	showVersion := fs.Bool("version", false, "print version and exit")
	fs.Usage = func() {
		fmt.Fprintf(stderr, "Usage: counterctl --address me@example.com | --domain example.com [--date YYYY-MM-DD] [--in|--out] [--json] [--config PATH] [--url URL]\n\n")
		fs.PrintDefaults()
	}
	if err := fs.Parse(args); err != nil {
		return 1
	}
	if *showVersion {
		fmt.Fprintf(stdout, "counterctl %s (commit %s, built %s)\n", version, commit, date)
		return 0
	}
	if (*address == "") == (*domain == "") {
		fmt.Fprintln(stderr, "exactly one of --address or --domain is required")
		fs.Usage()
		return 1
	}

	q := url.Values{}
	if *address != "" {
		addr, err := api.NormalizeEmail(*address)
		if err != nil {
			fmt.Fprintf(stderr, "invalid address %q: %v\n", *address, err)
			return 1
		}
		q.Set("address", addr)
	} else {
		q.Set("domain", strings.ToLower(strings.TrimSpace(*domain)))
	}
	if *dateFlag != "" {
		day, err := api.ParseDate(*dateFlag, now())
		if err != nil {
			fmt.Fprintln(stderr, err)
			return 1
		}
		q.Set("date", day)
		if t, _ := time.Parse("2006-01-02", day); now().UTC().Sub(t) > 31*24*time.Hour {
			fmt.Fprintf(stderr, "warning: %s is more than 31 days ago; counter keys expire after 31 days and will read as 0\n", day)
		}
	}

	cfg, err := loadConfig(*configPath, *urlFlag != "")
	if err != nil {
		fmt.Fprintln(stderr, err)
		return 1
	}
	token := cfg.AuthToken
	if v := getenv(config.EnvAuthToken); v != "" {
		token = v
	}
	base := *urlFlag
	if base == "" {
		base = baseURL(cfg.Addr)
		if host := strings.TrimPrefix(strings.TrimPrefix(base, "http://"), "https://"); !strings.HasPrefix(host, "localhost") && !strings.HasPrefix(host, "127.0.0.1") && !strings.HasPrefix(host, "[::1]") {
			fmt.Fprintf(stderr, "warning: %s is not local; the auth token travels over plaintext HTTP\n", base)
		}
	}

	req, err := http.NewRequest(http.MethodGet, base+"/counters?"+q.Encode(), nil)
	if err != nil {
		fmt.Fprintf(stderr, "bad URL: %v\n", err)
		return 1
	}
	if token != "" {
		req.Header.Set("Authorization", "Bearer "+token)
	}
	resp, err := (&http.Client{Timeout: 10 * time.Second}).Do(req)
	if err != nil {
		fmt.Fprintf(stderr, "request failed: %v\n", err)
		return 1
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		fmt.Fprintf(stderr, "reading response: %v\n", err)
		return 1
	}
	if resp.StatusCode != http.StatusOK {
		msg := strings.TrimSpace(string(body))
		var e api.ErrorResponse
		if json.Unmarshal(body, &e) == nil && e.Error != "" {
			msg = e.Error
		}
		fmt.Fprintf(stderr, "server answered %d: %s\n", resp.StatusCode, msg)
		return 1
	}
	if *jsonOut {
		stdout.Write(body)
		if len(body) == 0 || body[len(body)-1] != '\n' {
			io.WriteString(stdout, "\n")
		}
		return 0
	}
	var r api.CountersResponse
	if err := json.Unmarshal(body, &r); err != nil {
		fmt.Fprintf(stderr, "unexpected response: %v\n%s\n", err, body)
		return 1
	}
	showIn, showOut := flows(*inOnly, *outOnly)
	printTable(stdout, r, showIn, showOut)
	return 0
}

// flows picks the table columns: --in or --out alone restrict the table to
// that flow; neither or both show both flows.
func flows(in, out bool) (showIn, showOut bool) {
	if in == out {
		return true, true
	}
	return in, out
}

// loadConfig reads addr/auth_token from the daemon config. Without -config
// the default paths are tried; when none exists the CLI can still run with
// -url and the token from the environment.
func loadConfig(path string, haveURL bool) (cliConfig, error) {
	var cfg cliConfig
	if path == "" {
		found, err := config.Discover("")
		if err != nil {
			if haveURL {
				return cfg, nil
			}
			return cfg, err
		}
		path = found
	}
	if _, err := toml.DecodeFile(path, &cfg); err != nil {
		return cfg, fmt.Errorf("read config %s: %v", path, err)
	}
	return cfg, nil
}

// baseURL turns the daemon's listen address into a URL: ":5555" →
// "http://localhost:5555"; a full URL is used as is.
func baseURL(addr string) string {
	if addr == "" {
		addr = ":5555"
	}
	if strings.HasPrefix(addr, "http://") || strings.HasPrefix(addr, "https://") {
		return strings.TrimRight(addr, "/")
	}
	if strings.HasPrefix(addr, ":") {
		addr = "localhost" + addr
	}
	return "http://" + addr
}

func printTable(w io.Writer, r api.CountersResponse, showIn, showOut bool) {
	dom, org := "-", "-"
	if r.Domain != "" {
		dom = r.Domain
	}
	if r.OrgID != nil {
		org = strconv.FormatInt(*r.OrgID, 10)
	}
	fmt.Fprintf(w, "%s   date %s (UTC)   domain %s   org %s\n", r.Query, r.Date, dom, org)

	tw := tabwriter.NewWriter(w, 0, 8, 2, ' ', 0)
	head := []string{"LEVEL", "ID"}
	if showIn {
		head = append(head, "IN", "LIMIT")
	}
	if showOut {
		head = append(head, "OUT", "LIMIT")
	}
	fmt.Fprintln(tw, strings.Join(head, "\t"))
	for _, l := range r.Levels {
		cols := []string{l.Level, l.ID}
		if showIn {
			cols = append(cols, count(l.In, l.InLimit), limit(l.InLimit))
		}
		if showOut {
			cols = append(cols, count(l.Out, l.OutLimit), limit(l.OutLimit))
		}
		fmt.Fprintln(tw, strings.Join(cols, "\t"))
	}
	tw.Flush()

	if !r.Managed {
		fmt.Fprintln(w, "note: domain is not in rspamd_domains_db (not counted)")
	}
	if r.Exempt {
		fmt.Fprintln(w, "note: address is exempt from counting")
	}
}

// count renders a value, marked with "!" once it reached its limit.
func count(n int64, lim *int64) string {
	s := strconv.FormatInt(n, 10)
	if lim != nil && n >= *lim {
		s += "!"
	}
	return s
}

func limit(lim *int64) string {
	if lim == nil {
		return "-"
	}
	return strconv.FormatInt(*lim, 10)
}