// Command authcheck is a reference mitmux plugin - an Autorize-style
// authorization checker - and, deliberately, a template: it speaks
// mitmux's plugin wire protocol directly (raw JSON over the control
// socket, see PLUGINS.md) rather than importing mitmux's own internal
// Go packages, the same way a plugin written in any other language
// would have to. That's not a style preference - it's what actually
// proves PLUGINS.md's documented protocol is sufficient on its own,
// rather than silently depending on Go-internal conveniences a
// non-Go plugin author wouldn't have access to.
//
// What it does: for every captured request that carries an
// Authorization or Cookie header, resends the exact same request with
// that header stripped and compares the result. A resend that still
// succeeds where the original also succeeded means the endpoint
// doesn't actually enforce the authentication it appears to require -
// a missing-function-level-access-control bug, one of the more common
// real findings this class of check turns up. Matches are tagged
// "authcheck:bypass" on the original entry, with structured detail (the
// original status vs. the anonymous resend's) for the TUI's tag panel.
//
// This is the simplified half of what Burp's Autorize does: Autorize
// additionally supports swapping in a SECOND, lower-privileged
// identity's session and comparing against that - useful for catching
// cross-account IDORs a fully-anonymous check can't see. That needs a
// second credential as input, which this reference version doesn't
// take; a -low-priv-cookie flag doing that is a natural, small
// extension of the same pattern used here.
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"path/filepath"
)

// request/response mirror internal/ipc's wire structs field-for-field
// (see PLUGINS.md) - defined fresh here, not imported, so this file
// only ever exercises what's actually documented as the public
// protocol.
type request struct {
	Type      string `json:"type"`
	ID        int64  `json:"id,omitempty"`
	Scheme    string `json:"scheme,omitempty"`
	Host      string `json:"host,omitempty"`
	Raw       []byte `json:"raw,omitempty"`
	TagPlugin string `json:"tag_plugin,omitempty"`
	Tag       string `json:"tag,omitempty"`
	TagData   string `json:"tag_data,omitempty"`
}

type summary struct {
	ID     int64  `json:"id"`
	Method string `json:"method"`
	Scheme string `json:"scheme"`
	Host   string `json:"host"`
	Path   string `json:"path"`
	Source string `json:"source"`
}

type entryDetail struct {
	summary
	StatusCode int    `json:"status_code"`
	RequestRaw []byte `json:"request_raw"`
}

type response struct {
	Type    string       `json:"type"`
	Entries []summary    `json:"entries,omitempty"`
	New     *summary     `json:"new,omitempty"`
	Detail  *entryDetail `json:"detail,omitempty"`
	TagID   int64        `json:"tag_id,omitempty"`
	Error   string       `json:"error,omitempty"`
}

// client is a minimal request/response connection - send one request,
// read back one response, repeat. A plugin also needs a second,
// separate connection for "subscribe" (see main): that one is only
// ever written to once and then just read from continuously, so it
// doesn't need this type's request/response pairing at all.
type client struct {
	conn net.Conn
	enc  *json.Encoder
	dec  *json.Decoder
}

func dial(path string) (*client, error) {
	conn, err := net.Dial("unix", path)
	if err != nil {
		return nil, err
	}
	return &client{conn: conn, enc: json.NewEncoder(conn), dec: json.NewDecoder(conn)}, nil
}

func (c *client) call(req request) (response, error) {
	if err := c.enc.Encode(req); err != nil {
		return response{}, err
	}
	var resp response
	if err := c.dec.Decode(&resp); err != nil {
		return response{}, err
	}
	if resp.Type == "error" {
		return response{}, fmt.Errorf("%s", resp.Error)
	}
	return resp, nil
}

func defaultSocketPath() string {
	if rt := os.Getenv("XDG_RUNTIME_DIR"); rt != "" {
		return filepath.Join(rt, "mitmux.sock")
	}
	// Matches internal/ca.Dir() without importing it - see this file's
	// package doc for why plugins shouldn't reach into mitmux's own Go
	// internals even when it would be more convenient.
	cfg, err := os.UserConfigDir()
	if err != nil {
		return "mitmux.sock"
	}
	return filepath.Join(cfg, "mitmux", "mitmux.sock")
}

// authHeaders are checked in order; the first one present on a request
// is what gets stripped for the anonymous resend. Checking more than
// one matters: an API might authenticate via a bearer token while a
// browser-driven flow on the same host uses a session cookie, and both
// are worth checking independently rather than only ever picking one.
var authHeaders = []string{"Authorization", "Cookie"}

type result struct {
	OriginalStatus int    `json:"original_status"`
	ResendStatus   int    `json:"resend_status"`
	StrippedHeader string `json:"stripped_header"`
	Verdict        string `json:"verdict"`
}

func main() {
	socketPath := flag.String("socket", "", "daemon control socket path (default: same as mitmux itself)")
	pluginName := flag.String("name", "authcheck", "name this plugin tags entries as")
	flag.Parse()

	path := *socketPath
	if path == "" {
		path = defaultSocketPath()
	}

	actor, err := dial(path)
	if err != nil {
		log.Fatalf("dial %s: %v", path, err)
	}
	defer actor.conn.Close()

	subConn, err := net.Dial("unix", path)
	if err != nil {
		log.Fatalf("dial %s (subscribe): %v", path, err)
	}
	defer subConn.Close()
	if err := json.NewEncoder(subConn).Encode(request{Type: "subscribe"}); err != nil {
		log.Fatalf("subscribe: %v", err)
	}

	log.Printf("authcheck: watching live traffic on %s", path)
	dec := json.NewDecoder(subConn)
	for {
		var resp response
		if err := dec.Decode(&resp); err != nil {
			log.Fatalf("subscribe feed closed: %v", err)
		}
		if resp.Type != "new" || resp.New == nil {
			continue
		}
		sum := *resp.New
		// Never touch our own resends - the daemon records a Repeat()
		// as source="repeater", and reprocessing it would misfile the
		// very control-group requests this check depends on (it also
		// wouldn't loop: a resend already has its auth header removed,
		// so it would never match authHeaders again - but it's still
		// pointless work and pointless noise to try).
		if sum.Source != "proxy" {
			continue
		}
		if err := checkEntry(actor, *pluginName, sum.ID); err != nil {
			log.Printf("entry #%d: %v", sum.ID, err)
		}
	}
}

func checkEntry(c *client, pluginName string, id int64) error {
	resp, err := c.call(request{Type: "get", ID: id})
	if err != nil {
		return fmt.Errorf("get: %w", err)
	}
	if resp.Detail == nil {
		return fmt.Errorf("get: no detail in response")
	}
	detail := *resp.Detail

	req, err := http.ReadRequest(bufio.NewReader(bytes.NewReader(detail.RequestRaw)))
	if err != nil {
		return nil // not a well-formed request we can safely reparse - skip, not fatal
	}

	var strippedHeader string
	for _, h := range authHeaders {
		if req.Header.Get(h) != "" {
			strippedHeader = h
			break
		}
	}
	if strippedHeader == "" {
		return nil // nothing to check
	}

	req2, err := http.ReadRequest(bufio.NewReader(bytes.NewReader(detail.RequestRaw)))
	if err != nil {
		return nil
	}
	req2.Header.Del(strippedHeader)

	var buf bytes.Buffer
	if err := req2.Write(&buf); err != nil {
		return fmt.Errorf("rebuild request: %w", err)
	}

	repeatResp, err := c.call(request{Type: "repeat", Scheme: detail.Scheme, Host: detail.Host, Raw: buf.Bytes()})
	if err != nil {
		return fmt.Errorf("repeat: %w", err)
	}
	if repeatResp.Detail == nil {
		return fmt.Errorf("repeat: no detail in response")
	}
	resentStatus := repeatResp.Detail.StatusCode

	suspicious := successClass(detail.StatusCode) && successClass(resentStatus)
	if !suspicious {
		return nil // matches Burp's own default: only surface likely findings, not every check performed
	}

	data, err := json.Marshal(result{
		OriginalStatus: detail.StatusCode,
		ResendStatus:   resentStatus,
		StrippedHeader: strippedHeader,
		Verdict:        "bypass",
	})
	if err != nil {
		return fmt.Errorf("marshal result: %w", err)
	}

	if _, err := c.call(request{Type: "tag_entry", ID: id, TagPlugin: pluginName, Tag: "authcheck:bypass", TagData: string(data)}); err != nil {
		return fmt.Errorf("tag_entry: %w", err)
	}
	log.Printf("#%d %s %s -> possible auth bypass: %d without %s still got %d",
		id, detail.Method, detail.Path, detail.StatusCode, strippedHeader, resentStatus)
	return nil
}

func successClass(status int) bool {
	return status >= 200 && status < 400
}