// Package proxy реализует форвардный HTTP(S)-прокси с перехватом
// трафика (MITM), аналог Burp Proxy / mitmproxy. Браузер настраивается
// ходить через Server.Addr; для HTTP запросы читаются напрямую, для
// HTTPS используется CONNECT-туннель с подменой TLS-сертификата на
// лету (см. CA.LeafFor) — отсюда необходимость один раз импортировать
// ca.pem в доверенные корни браузера/ОС.
//
// Известные упрощения первой версии (не блокирующие базовый сценарий
// перехвата и правки запроса, но стоящие в очереди на доработку):
//   - hop-by-hop заголовки (Connection, Proxy-Connection и т.п.) не
//     вырезаются перед проксированием;
//   - WebSocket-апгрейды не поддерживаются;
//   - HTTP/2 на туннеле не форсируется (сервер объявляет только h1
//     в ALPN, так что браузер сам откатится на HTTP/1.1 для
//     туннелированных доменов).
package proxy

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"sync"
	"time"

	"github.com/r3g1tpr0cs/burterm/internal/sitemap"
)

// maxResponseCapture — сколько байт тела ответа сохраняется в
// Entry.ResponseBody для полнотекстового поиска в Project. Больше, чем
// в Intruder (4КБ) — трафик через Proxy обычно не такой массовый, можно
// позволить себе больше на хранение.
const maxResponseCapture = 8 * 1024

// Server — форвардный прокси. Каждый прошедший через него запрос
// попадает в Store, откуда его можно подхватить в Request Builder.
type Server struct {
	Addr  string
	CA    *CA
	Store *Store

	// scope, если задан, ограничивает то, что реально попадает в Store —
	// сам прокси всё равно обслуживает ЛЮБОЙ трафик (клиент получает
	// ответ независимо от scope), просто то, что вне scope, не засоряет
	// историю. nil (значение по умолчанию) — логировать всё, как раньше.
	// Под мьютексом: SetScope дёргается из горутины TUI, а capture/
	// captureErr читают его из совсем других горутин — по одной на
	// соединение.
	scopeMu sync.RWMutex
	scope   *sitemap.Scope

	client *http.Client
}

// SetScope задаёт (или снимает, если nil) scope, ограничивающий то, что
// попадает в Store. Потокобезопасно относительно capture/captureErr.
// Передаётся независимая копия сверху (см. TUI-слой) — Server не должен
// делить один и тот же *sitemap.Scope с чем-то ещё мутируемым конкурентно,
// иначе Matches() читал бы срез правил без синхронизации с их же записью.
func (s *Server) SetScope(scope *sitemap.Scope) {
	s.scopeMu.Lock()
	s.scope = scope
	s.scopeMu.Unlock()
}

// inScope — под nil-scope пропускает всё (обратная совместимость с
// поведением до появления этой фичи).
func (s *Server) inScope(url string) bool {
	s.scopeMu.RLock()
	scope := s.scope
	s.scopeMu.RUnlock()
	if scope == nil {
		return true
	}
	return scope.Matches(url)
}

// New создаёт Server, готовый к ListenAndServe.
func New(addr string, ca *CA, store *Store) *Server {
	return &Server{
		Addr:  addr,
		CA:    ca,
		Store: store,
		client: &http.Client{
			// Собственный Transport — чтобы не поймать петлю через
			// переменные окружения HTTP_PROXY/HTTPS_PROXY.
			Transport: &http.Transport{Proxy: nil},
			// Редиректы не разворачиваем автоматически — пользователь
			// должен увидеть 30x как отдельную перехваченную транзакцию,
			// а не только её конечную точку.
			CheckRedirect: func(*http.Request, []*http.Request) error {
				return http.ErrUseLastResponse
			},
		},
	}
}

// ListenAndServe запускает прокси и блокируется, пока ctx не отменён
// или не произойдёт фатальная ошибка сервера.
func (s *Server) ListenAndServe(ctx context.Context) error {
	srv := &http.Server{
		Addr:    s.Addr,
		Handler: http.HandlerFunc(s.serve),
	}

	errCh := make(chan error, 1)
	go func() { errCh <- srv.ListenAndServe() }()

	select {
	case <-ctx.Done():
		shutdownCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
		defer cancel()
		return srv.Shutdown(shutdownCtx)
	case err := <-errCh:
		return err
	}
}

func (s *Server) serve(w http.ResponseWriter, r *http.Request) {
	if r.Method == http.MethodConnect {
		s.serveConnect(w, r)
		return
	}
	s.servePlain(w, r)
}

// servePlain обрабатывает обычный нешифрованный прокси-запрос — браузер
// шлёт его с абсолютным URI в request line (GET http://host/path HTTP/1.1).
func (s *Server) servePlain(w http.ResponseWriter, r *http.Request) {
	// DumpRequest вычитывает r.Body и сам восстанавливает его новым
	// io.ReadCloser с теми же байтами — r остаётся пригоден для
	// дальнейшей пересылки.
	dump, _ := httputil.DumpRequest(r, true)

	outURL := *r.URL
	if outURL.Scheme == "" {
		outURL.Scheme = "http"
	}
	if outURL.Host == "" {
		outURL.Host = r.Host
	}

	outReq, err := http.NewRequest(r.Method, outURL.String(), r.Body)
	if err != nil {
		s.captureErr(r, err)
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	outReq.Header = r.Header.Clone()

	resp, err := s.client.Do(outReq)
	if err != nil {
		s.captureErr(r, err)
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()

	// Читаем тело целиком — не обрезаем то, что реально уходит клиенту.
	// Ошибку чтения намеренно не считаем фатальной: то, что успели
	// прочитать, всё равно форвардим и сохраняем как есть.
	fullBody, _ := io.ReadAll(resp.Body)

	s.capture(r, dump, resp, fullBody)

	for k, vv := range resp.Header {
		for _, v := range vv {
			w.Header().Add(k, v)
		}
	}
	w.WriteHeader(resp.StatusCode)
	w.Write(fullBody)
}

// serveConnect обрабатывает CONNECT (HTTPS-туннель): хайджекает
// соединение, отвечает "200 Connection Established", затем сам
// выступает TLS-сервером с сертификатом, подписанным локальным CA
// (GetCertificate генерирует лист по SNI на лету), и внутри
// расшифрованного канала читает и пересылает HTTP-запросы как обычный
// клиент — это и есть MITM-перехват.
func (s *Server) serveConnect(w http.ResponseWriter, r *http.Request) {
	hijacker, ok := w.(http.Hijacker)
	if !ok {
		http.Error(w, "хайджек соединения не поддерживается", http.StatusInternalServerError)
		return
	}
	clientConn, _, err := hijacker.Hijack()
	if err != nil {
		return
	}
	defer clientConn.Close()

	host, _, err := net.SplitHostPort(r.Host)
	if err != nil {
		host = r.Host // Host уже без порта, если SplitHostPort не смог разобрать
	}

	if _, err := clientConn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n")); err != nil {
		return
	}

	tlsConn := tls.Server(clientConn, &tls.Config{
		GetCertificate: func(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
			name := hello.ServerName
			if name == "" {
				name = host
			}
			return s.CA.LeafFor(name)
		},
	})
	defer tlsConn.Close()

	if err := tlsConn.Handshake(); err != nil {
		// Обычно значит, что браузер ещё не доверяет нашему CA —
		// это ожидаемо до первого импорта ca.pem, молча выходим.
		return
	}

	br := bufio.NewReader(tlsConn)
	for {
		req, err := http.ReadRequest(br)
		if err != nil {
			return
		}

		// http.ReadRequest для relative-form request line оставляет
		// req.URL.Scheme/Host пустыми — достраиваем их из адреса
		// CONNECT-туннеля, иначе пересылать запрос некуда.
		req.URL.Scheme = "https"
		req.URL.Host = r.Host
		req.RequestURI = ""

		dump, _ := httputil.DumpRequest(req, true)

		resp, err := s.client.Do(req)
		if err != nil {
			s.captureErr(req, err)
			return
		}

		// resp.Write() ниже сам сериализует тело из resp.Body — а нам
		// то же тело нужно ещё и для capture(). Тело можно прочитать
		// только один раз, поэтому читаем целиком сами и подменяем
		// resp.Body свежим io.NopCloser с теми же байтами, прежде чем
		// отдать resp дальше — тот же приём, что httputil.DumpRequest
		// делает для тела запроса.
		fullBody, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		resp.Body = io.NopCloser(bytes.NewReader(fullBody))

		s.capture(req, dump, resp, fullBody)

		writeErr := resp.Write(tlsConn)
		if writeErr != nil {
			return
		}
	}
}

func (s *Server) capture(r *http.Request, raw []byte, resp *http.Response, fullBody []byte) {
	url := r.URL.String()
	if !s.inScope(url) {
		return
	}

	captured := fullBody
	if len(captured) > maxResponseCapture {
		captured = captured[:maxResponseCapture]
	}
	s.Store.Add(Entry{
		Timestamp:    time.Now(),
		Method:       r.Method,
		URL:          url,
		Raw:          string(raw),
		StatusCode:   resp.StatusCode,
		Length:       int64(len(fullBody)), // фактически прочитанные байты, не ненадёжный resp.ContentLength (часто -1 при chunked)
		Headers:      resp.Header,
		ResponseBody: string(captured),
	})
}

func (s *Server) captureErr(r *http.Request, err error) {
	url := ""
	if r.URL != nil {
		url = r.URL.String()
	}
	if url != "" && !s.inScope(url) {
		return
	}
	s.Store.Add(Entry{
		Timestamp: time.Now(),
		Method:    r.Method,
		URL:       url,
		Err:       fmt.Errorf("проксирование не удалось: %w", err),
	})
}