// Command bpscanner is a reference mitmux plugin - a Backslash Powered
// Scanner-style generic injection detector. Where most scanners test
// known payloads for known technologies (`' OR 1=1--` for SQL, `<script>`
// for XSS), this tests a more general question: does the backend do
// anything different with syntactically-significant characters than it
// does with an equal-length string of inert ones? That question doesn't
// need to know what the backend is built on to be worth asking, which
// is the appeal of the real tool this borrows its name and idea from.
//
// For every existing query parameter on a captured GET request, sends
// two same-length replacement values wrapped in a stable marker: one
// full of characters that mean something in a lot of common contexts
// (quotes, angle brackets, shell/template metacharacters - `'"\<>(){}$;|&`),
// one full of harmless filler. If the marker comes back intact for the
// filler value but broken, altered, or missing for the special-character
// one - or the two get different status codes outright - something
// downstream is interpreting those characters rather than treating the
// parameter as inert data, tagged "bpscanner:hit" complementing whatever
// paramminer already covers (which finds parameters that shouldn't
// exist at all; this tests parameters that already do).
//
// Speaks the wire protocol directly, no internal/ipc import - see
// plugins/authcheck's package doc for why, and PLUGINS.md for the
// protocol.
package main

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

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"`
	ResponseRaw []byte `json:"response_raw"`
}

type response struct {
	Type   string       `json:"type"`
	New    *summary     `json:"new,omitempty"`
	Detail *entryDetail `json:"detail,omitempty"`
	Error  string       `json:"error,omitempty"`
}

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")
	}
	cfg, err := os.UserConfigDir()
	if err != nil {
		return "mitmux.sock"
	}
	return filepath.Join(cfg, "mitmux", "mitmux.sock")
}

// marker wraps both the special and control payloads so a response can
// be checked for whether it survived intact, not just whether the
// response as a whole "looks different" - a much more targeted signal
// than length/status diffing alone. Chosen unlikely to collide with
// real content.
const marker = "zzMARKzz"

// special covers characters meaningful across a lot of common
// contexts at once - SQL quoting, HTML/JS, shell metacharacters,
// template/expression syntax - without committing to any one of them.
// filler is the same length, entirely inert.
const (
	special = `'"\<>(){}$;|&`
	filler  = `AAAAAAAAAAAAA`
)

type hit struct {
	Parameter     string `json:"parameter"`
	ControlStatus int    `json:"control_status"`
	SpecialStatus int    `json:"special_status"`
	ControlIntact bool   `json:"control_marker_intact"`
	SpecialIntact bool   `json:"special_marker_intact"`
}

func main() {
	socketPath := flag.String("socket", "", "daemon control socket path (default: same as mitmux itself)")
	pluginName := flag.String("name", "bpscanner", "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)
	}

	// Dedup by endpoint AND parameter name - unlike paramminer (which
	// only cares about the endpoint, since it's testing candidates that
	// don't depend on what's already there), this needs a fresh entry
	// per distinct parameter worth testing.
	probed := map[string]bool{}

	log.Printf("bpscanner: 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
		if sum.Source != "proxy" || sum.Method != "GET" {
			continue
		}
		if err := scanEntry(actor, *pluginName, sum.ID, probed); err != nil {
			log.Printf("entry #%d: %v", sum.ID, err)
		}
	}
}

func scanEntry(c *client, pluginName string, id int64, probed map[string]bool) 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

	baseReq, 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
	}
	query := baseReq.URL.Query()
	if len(query) == 0 {
		return nil // nothing to mutate
	}

	var hits []hit
	for param := range query {
		key := detail.Method + " " + detail.Scheme + "://" + detail.Host + detail.Path + "?" + param
		if probed[key] {
			continue
		}
		probed[key] = true

		controlStatus, controlIntact, err := probeParam(c, detail.Scheme, detail.Host, baseReq, param, filler)
		if err != nil {
			log.Printf("entry #%d param %q control probe: %v", id, param, err)
			continue
		}
		specialStatus, specialIntact, err := probeParam(c, detail.Scheme, detail.Host, baseReq, param, special)
		if err != nil {
			log.Printf("entry #%d param %q special probe: %v", id, param, err)
			continue
		}

		// The interesting case: the harmless filler survives intact but
		// the special-character payload doesn't, or the two get
		// different status codes outright. Both marker-intact and both
		// broken (or both status codes matching) isn't a finding - that
		// just means the value never reaches anywhere meaningful.
		if controlStatus != specialStatus || (controlIntact && !specialIntact) {
			hits = append(hits, hit{
				Parameter:     param,
				ControlStatus: controlStatus,
				SpecialStatus: specialStatus,
				ControlIntact: controlIntact,
				SpecialIntact: specialIntact,
			})
		}
	}

	if len(hits) == 0 {
		return nil
	}

	data, err := json.Marshal(hits)
	if err != nil {
		return fmt.Errorf("marshal hits: %w", err)
	}
	if _, err := c.call(request{Type: "tag_entry", ID: id, TagPlugin: pluginName, Tag: "bpscanner:hit", TagData: string(data)}); err != nil {
		return fmt.Errorf("tag_entry: %w", err)
	}
	names := make([]string, len(hits))
	for i, h := range hits {
		names[i] = h.Parameter
	}
	log.Printf("#%d %s -> possible injection point(s): %s", id, detail.Path, strings.Join(names, ", "))
	return nil
}

// probeParam resends baseReq with param's value replaced by
// marker+payload+marker, returning the response's status code and
// whether both marker occurrences came back intact and unmodified.
func probeParam(c *client, scheme, host string, baseReq *http.Request, param, payload string) (status int, markerIntact bool, err error) {
	u := *baseReq.URL
	q := u.Query()
	value := marker + payload + marker
	q.Set(param, value)
	u.RawQuery = q.Encode()

	req2 := baseReq.Clone(baseReq.Context())
	req2.URL = &u

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

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

	body := string(resp.Detail.ResponseRaw)
	intact := strings.Count(body, marker) >= 2 && strings.Contains(body, marker+payload+marker)
	return resp.Detail.StatusCode, intact, nil
}