upd

internal/proxy/server.go

@@ -0,0 +1,87 @@
+package proxy
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"log"
+	"net/http"
+	"os"
+	"strings"
+
+	"github.com/elazarl/goproxy"
+
+	"git.gulenok.ru/greenhaze/gaslight/internal/capture"
+	"git.gulenok.ru/greenhaze/gaslight/internal/config"
+	"git.gulenok.ru/greenhaze/gaslight/internal/domain"
+	"git.gulenok.ru/greenhaze/gaslight/internal/rewrite"
+	"git.gulenok.ru/greenhaze/gaslight/internal/tlsutil"
+)
+
+func Run(cfg config.Config) error {
+	if err := os.MkdirAll(cfg.OutputDir, 0o755); err != nil {
+		return fmt.Errorf("failed to create output directory: %w", err)
+	}
+
+	caCert, err := tlsutil.LoadCACert(cfg.CertPath, cfg.KeyPath)
+	if err != nil {
+		return fmt.Errorf("failed to load CA certificate/key: %w", err)
+	}
+	goproxy.GoproxyCa = *caCert
+
+	imageSaver := capture.NewImageSaver(cfg.OutputDir)
+	rewriter := rewrite.NewProcessor(rewrite.DefaultQuestionTestHook)
+
+	p := goproxy.NewProxyHttpServer()
+	p.Verbose = false
+	p.OnRequest().HandleConnect(goproxy.AlwaysMitm)
+	p.OnResponse().DoFunc(func(resp *http.Response, _ *goproxy.ProxyCtx) *http.Response {
+		if resp == nil || resp.Request == nil || resp.Body == nil {
+			return resp
+		}
+
+		if !domain.IsMckoDomain(resp.Request.Host) {
+			return resp
+		}
+
+		original, err := io.ReadAll(resp.Body)
+		_ = resp.Body.Close()
+		if err != nil {
+			log.Printf("read body failed for %s: %v", resp.Request.URL.String(), err)
+			return resp
+		}
+
+		contentType := strings.ToLower(resp.Header.Get("Content-Type"))
+		if capture.IsImageResponse(contentType, resp.Request.URL.Path) {
+			if err := imageSaver.Save(resp, original); err != nil {
+				log.Printf("save image failed for %s: %v", resp.Request.URL.String(), err)
+			}
+			setResponseBody(resp, original)
+			return resp
+		}
+
+		if !rewrite.IsHTMLResponse(contentType) {
+			setResponseBody(resp, original)
+			return resp
+		}
+
+		modified, err := rewriter.RewriteIfNeeded(resp.Request, resp.Header.Get("Content-Encoding"), original)
+		if err != nil {
+			log.Printf("html rewrite failed for %s: %v", resp.Request.URL.String(), err)
+			setResponseBody(resp, original)
+			return resp
+		}
+
+		setResponseBody(resp, modified)
+		return resp
+	})
+
+	log.Printf("proxy listening on %s", cfg.ListenAddr)
+	return http.ListenAndServe(cfg.ListenAddr, p)
+}
+
+func setResponseBody(resp *http.Response, body []byte) {
+	resp.Body = io.NopCloser(bytes.NewReader(body))
+	resp.ContentLength = int64(len(body))
+	resp.Header.Set("Content-Length", fmt.Sprint(len(body)))
+}