// ipabusedb-docker-ssh reads the logs of a dockerised sshd via
// "docker logs -t", extracts brute-force / invalid-user attempts, and
// emits a CSV suitable for the AbuseIPDB bulk-report endpoint.
//
// Environment:
//
//	DOCKER_CONTAINER_NAME  container whose logs are read (required)
//	IPABUSEDB_API_KEY      AbuseIPDB API key (required with --submit)
//
// Flags:
//
//	--since 24h   only consider log lines newer than now()-duration
//	-o FILE       write CSV to FILE instead of stdout
//	--input FILE  parse FILE (or "-" for stdin) instead of docker logs
//	--submit      POST the CSV to the AbuseIPDB bulk-report endpoint
package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/csv"
	"errors"
	"flag"
	"fmt"
	"io"
	"log/slog"
	"mime/multipart"
	"net/http"
	"net/netip"
	"os"
	"os/exec"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"
)

const (
	bulkReportURL = "https://api.abuseipdb.com/api/v2/bulk-report"
	// AbuseIPDB categories: 18 = Brute-Force, 22 = SSH.
	categories   = "18,22"
	excludedUser = "git"
	// AbuseIPDB rejects bulk reports with more than this many rows.
	bulkReportMaxRows = 10000
)

// patterns match sshd messages that evidence an unauthorised login
// attempt. Group 1 is the username, group 2 the source address. Order
// is priority: the earliest matching pattern supplies the report comment.
var patterns = []*regexp.Regexp{
	regexp.MustCompile(`Invalid user (\S+) from (\S+) port \d+`),
	regexp.MustCompile(`User (\S+) from (\S+) not allowed because not listed in AllowUsers`),
	regexp.MustCompile(`Failed password for invalid user (\S+) from (\S+) port \d+`),
	regexp.MustCompile(`Connection closed by invalid user (\S+) (\S+) port \d+`),
	regexp.MustCompile(`Disconnected from invalid user (\S+) (\S+) port \d+`),
}

type event struct {
	ip      netip.Addr
	user    string
	comment string
	prio    int // index into patterns; lower is more descriptive
	when    time.Time
	count   int
}

// parser accumulates one event per source IP, keeping the most recent
// matching message and counting how many matched in total.
type parser struct {
	cutoff time.Time

	mu      sync.Mutex
	events  map[netip.Addr]*event
	skipped int
}

func newParser(cutoff time.Time) *parser {
	return &parser{cutoff: cutoff, events: map[netip.Addr]*event{}}
}

// handleLine parses one "docker logs -t" line: an RFC3339Nano timestamp,
// a space, then the message.
func (p *parser) handleLine(line string) {
	line = strings.TrimRight(line, "\r\n")
	ts, msg, ok := strings.Cut(line, " ")
	if !ok {
		return
	}
	when, err := time.Parse(time.RFC3339Nano, ts)
	if err != nil {
		p.mu.Lock()
		p.skipped++
		p.mu.Unlock()
		return
	}
	if !p.cutoff.IsZero() && when.Before(p.cutoff) {
		return
	}
	for prio, re := range patterns {
		m := re.FindStringSubmatch(msg)
		if m == nil {
			continue
		}
		user, addr := m[1], m[2]
		if user == excludedUser {
			return
		}
		ip, err := netip.ParseAddr(addr)
		if err != nil || !reportable(ip) {
			return
		}
		p.record(ip, user, strings.TrimSpace(m[0]), prio, when)
		return
	}
}

func reportable(ip netip.Addr) bool {
	return ip.IsGlobalUnicast() && !ip.IsPrivate()
}

// record merges an observation into the per-IP event: the report date is
// the latest observation, the comment comes from the most descriptive
// pattern (and, within the same pattern, the latest observation).
func (p *parser) record(ip netip.Addr, user, comment string, prio int, when time.Time) {
	p.mu.Lock()
	defer p.mu.Unlock()
	ev, ok := p.events[ip]
	if !ok {
		p.events[ip] = &event{ip: ip, user: user, comment: comment, prio: prio, when: when, count: 1}
		return
	}
	ev.count++
	newer := when.After(ev.when)
	if newer {
		ev.when = when
	}
	if prio < ev.prio || (prio == ev.prio && newer) {
		ev.prio, ev.user, ev.comment = prio, user, comment
	}
}

func (p *parser) ingest(r io.Reader) error {
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
	for sc.Scan() {
		p.handleLine(sc.Text())
	}
	return sc.Err()
}

func (p *parser) sorted() []*event {
	p.mu.Lock()
	defer p.mu.Unlock()
	out := make([]*event, 0, len(p.events))
	for _, ev := range p.events {
		out = append(out, ev)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].when.Before(out[j].when) })
	return out
}

func writeCSV(w io.Writer, events []*event) error {
	cw := csv.NewWriter(w)
	if err := cw.Write([]string{"IP", "Categories", "Comment", "ReportDate"}); err != nil {
		return err
	}
	for _, ev := range events {
		comment := ev.comment
		if ev.count > 1 {
			comment = fmt.Sprintf("%s (%d events)", ev.comment, ev.count)
		}
		row := []string{
			ev.ip.String(),
			categories,
			comment,
			ev.when.UTC().Format(time.RFC3339),
		}
		if err := cw.Write(row); err != nil {
			return err
		}
	}
	cw.Flush()
	return cw.Error()
}

// readDockerLogs runs "docker logs -t [--since D] NAME" and feeds both
// stdout and stderr (sshd logs to stderr) through the parser.
func readDockerLogs(ctx context.Context, p *parser, container string, since time.Duration) error {
	args := []string{"logs", "-t"}
	if since > 0 {
		args = append(args, "--since", since.String())
	}
	args = append(args, container)
	cmd := exec.CommandContext(ctx, "docker", args...)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}
	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}
	if err := cmd.Start(); err != nil {
		return fmt.Errorf("starting docker: %w", err)
	}

	var wg sync.WaitGroup
	errs := make(chan error, 2)
	for _, r := range []io.Reader{stdout, stderr} {
		wg.Add(1)
		go func(r io.Reader) {
			defer wg.Done()
			if err := p.ingest(r); err != nil {
				errs <- err
			}
		}(r)
	}
	wg.Wait()
	close(errs)
	if err := cmd.Wait(); err != nil {
		return fmt.Errorf("docker logs %s: %w", container, err)
	}
	for err := range errs {
		if err != nil {
			return err
		}
	}
	return nil
}

func readInput(p *parser, path string) error {
	if path == "-" {
		return p.ingest(os.Stdin)
	}
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	return p.ingest(f)
}

func submit(ctx context.Context, apiKey string, csvData []byte) ([]byte, error) {
	var body bytes.Buffer
	mw := multipart.NewWriter(&body)
	part, err := mw.CreateFormFile("csv", "reports.csv")
	if err != nil {
		return nil, err
	}
	if _, err := part.Write(csvData); err != nil {
		return nil, err
	}
	if err := mw.Close(); err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodPost, bulkReportURL, &body)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Key", apiKey)
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Content-Type", mw.FormDataContentType())

	client := &http.Client{Timeout: 2 * time.Minute}
	resp, err := client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	respBody, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode/100 != 2 {
		return respBody, fmt.Errorf("abuseipdb returned %s", resp.Status)
	}
	return respBody, nil
}

func run() error {
	since := flag.Duration("since", 0, "only report events newer than now()-duration (e.g. 24h)")
	outPath := flag.String("o", "", "write CSV to this file instead of stdout")
	inPath := flag.String("input", "", "parse this file (or - for stdin) instead of running docker logs")
	doSubmit := flag.Bool("submit", false, "POST the CSV to the AbuseIPDB bulk-report endpoint")
	flag.Parse()

	apiKey := os.Getenv("IPABUSEDB_API_KEY")
	container := os.Getenv("DOCKER_CONTAINER_NAME")
	if *doSubmit && apiKey == "" {
		return errors.New("IPABUSEDB_API_KEY must be set in the environment to --submit")
	}
	if *inPath == "" && container == "" {
		return errors.New("DOCKER_CONTAINER_NAME must be set in the environment")
	}

	var cutoff time.Time
	if *since > 0 {
		cutoff = time.Now().Add(-*since)
	}
	p := newParser(cutoff)

	ctx := context.Background()
	var err error
	if *inPath != "" {
		err = readInput(p, *inPath)
	} else {
		err = readDockerLogs(ctx, p, container, *since)
	}
	if err != nil {
		return err
	}

	events := p.sorted()
	if len(events) > bulkReportMaxRows {
		slog.Warn("truncating to AbuseIPDB bulk report limit", "limit", bulkReportMaxRows, "found", len(events))
		events = events[len(events)-bulkReportMaxRows:]
	}
	slog.Info("parsed logs", "abusers", len(events), "unparseable_lines", p.skipped)

	var buf bytes.Buffer
	if err := writeCSV(&buf, events); err != nil {
		return err
	}

	out := io.Writer(os.Stdout)
	if *outPath != "" {
		f, err := os.Create(*outPath)
		if err != nil {
			return err
		}
		defer f.Close()
		out = f
	}
	if _, err := out.Write(buf.Bytes()); err != nil {
		return err
	}

	if !*doSubmit {
		return nil
	}
	if len(events) == 0 {
		slog.Info("nothing to submit")
		return nil
	}
	resp, err := submit(ctx, apiKey, buf.Bytes())
	if len(resp) > 0 {
		fmt.Fprintln(os.Stderr, string(resp))
	}
	if err != nil {
		return err
	}
	slog.Info("submitted bulk report", "rows", len(events))
	return nil
}

func main() {
	slog.SetDefault(slog.New(slog.NewTextHandler(os.Stderr, nil)))
	if err := run(); err != nil {
		slog.Error(err.Error())
		os.Exit(1)
	}
}