added script support in aitm proxy for setting dynamic upstream proxy

Signed-off-by: RonniSkansing <rskansing@gmail.com>

backend/app/server.go

@@ -99,6 +99,7 @@ func NewServer(
 		services.IPAllowList,
 		repositories.Option,
 		services.Option,
+		services.Script,
 		trustedProxies,
 	)
 

backend/app/services.go

@@ -232,6 +232,7 @@ func NewServices(
 		CampaignTemplateService: campaignTemplate,
 		DomainService:           domain,
 		ProxySessionManager:     proxySessionManager,
+		ScriptEnabled:           scriptEnabled,
 	}
 	ipAllowListService := service.NewIPAllowListService(logger, repositories.Proxy)
 	email := &service.Email{

backend/proxy/proxy.go

@@ -125,6 +125,7 @@ type ProxyHandler struct {
 	IPAllowListService          *service.IPAllowListService
 	OptionRepository            *repository.Option
 	OptionService               *service.Option
+	ScriptService               *service.Script
 	cookieName                  string
 	trustedProxies              []string
 }
@@ -144,6 +145,7 @@ func NewProxyHandler(
 	ipAllowListService *service.IPAllowListService,
 	optionRepo *repository.Option,
 	optionService *service.Option,
+	scriptService *service.Script,
 	trustedProxies []string,
 ) *ProxyHandler {
 	// get proxy cookie name from database
@@ -167,6 +169,7 @@ func NewProxyHandler(
 		IPAllowListService:          ipAllowListService,
 		OptionRepository:            optionRepo,
 		OptionService:               optionService,
+		ScriptService:               scriptService,
 		cookieName:                  cookieName,
 		trustedProxies:              trustedProxies,
 	}
@@ -236,18 +239,22 @@ func (m *ProxyHandler) HandleHTTPRequest(w http.ResponseWriter, req *http.Reques
 	// this ensures custom user-agent replacements work with impersonation
 	m.applyEarlyRequestHeaderReplacements(req, reqCtx)
 
-	// create http client with optional browser impersonation
-	client, err := m.createHTTPClientWithImpersonation(req, reqCtx, reqCtx.ProxyConfig)
-	if err != nil {
-		return errors.Errorf("failed to create proxy HTTP client: %w", err)
-	}
-
-	// process request
+	// process request. this resolves or creates the session, so the session
+	// script has run and any upstream proxy it selected is set on the session
+	// before the client is built below.
 	modifiedReq, resp := m.processRequestWithContext(req, reqCtx)
 	if resp != nil {
 		return m.writeResponse(w, resp)
 	}
 
+	// create http client with optional browser impersonation. built after the
+	// session is resolved so it can use the upstream proxy the session script
+	// selected for this session.
+	client, err := m.createHTTPClientWithImpersonation(modifiedReq, reqCtx, reqCtx.ProxyConfig)
+	if err != nil {
+		return errors.Errorf("failed to create proxy HTTP client: %w", err)
+	}
+
 	// prepare request for target server
 	m.prepareRequestForTarget(modifiedReq, client, reqCtx.UsedImpersonation)
 
@@ -1455,6 +1462,11 @@ func (m *ProxyHandler) createNewSession(
 	// initialize session data
 	m.initializeSession(session, sessionConfig)
 
+	// run the session script once for this new session so it can pick the
+	// upstream proxy from the incoming connection. set before the session is
+	// stored so later requests read a stable value.
+	session.UpstreamProxy = m.runSessionScript(req, reqCtx)
+
 	// store session
 	m.SessionManager.StoreSession(session.ID, session)
 	if campaignRecipientID != nil {
@@ -1464,6 +1476,88 @@ func (m *ProxyHandler) createNewSession(
 	return session, nil
 }
 
+// runSessionScript runs the configured session script once for a new session
+// and returns the upstream proxy it selects. The script receives the incoming
+// connection (ip, country, asns, ja4, headers) and may return an object with a
+// "proxy" field. On any error, or when no script is configured, it returns an
+// empty string so the session falls back to the proxy from the yaml config.
+func (m *ProxyHandler) runSessionScript(req *http.Request, reqCtx *RequestContext) string {
+	if m.ScriptService == nil || reqCtx.ProxyConfig == nil {
+		return ""
+	}
+	scriptName := strings.TrimSpace(reqCtx.ProxyConfig.Script)
+	if scriptName == "" {
+		return ""
+	}
+
+	// resolve the company scope from the campaign so the lookup stays within
+	// the campaign's company plus global scripts
+	var companyID *uuid.UUID
+	if reqCtx.Campaign != nil {
+		if cid, err := reqCtx.Campaign.CompanyID.Get(); err == nil {
+			companyID = &cid
+		}
+	}
+
+	input := m.buildScriptRequestInput(req, reqCtx)
+
+	out, err := m.ScriptService.RunCallable(req.Context(), companyID, scriptName, input)
+	if err != nil {
+		m.logger.Errorw("session script failed, falling back to config proxy",
+			"script", scriptName,
+			"error", err,
+		)
+		return ""
+	}
+	if out == nil {
+		return ""
+	}
+	if proxy, ok := out["proxy"].(string); ok {
+		return strings.TrimSpace(proxy)
+	}
+	return ""
+}
+
+// buildScriptRequestInput captures the incoming connection for the session
+// script: the trusted proxy aware ip, its country and asns from the ipdata
+// store, the ja4 fingerprint, and the request headers.
+func (m *ProxyHandler) buildScriptRequestInput(req *http.Request, reqCtx *RequestContext) map[string]interface{} {
+	ip := utils.ExtractClientIP(req, m.trustedProxies)
+	if host, _, err := net.SplitHostPort(ip); err == nil {
+		ip = host
+	}
+
+	headers := map[string]interface{}{}
+	for k := range req.Header {
+		headers[strings.ToLower(k)] = req.Header.Get(k)
+	}
+
+	country := ""
+	asns := []map[string]interface{}{}
+	if store := ipdata.Get(); store != nil {
+		if c, ok := store.LookupCountry(ip); ok {
+			country = c
+		}
+		for _, a := range store.LookupASNDetails(ip) {
+			asns = append(asns, map[string]interface{}{
+				"number": a.ASN,
+				"name":   a.Name,
+			})
+		}
+	}
+
+	return map[string]interface{}{
+		"ip":             ip,
+		"country":        country,
+		"asns":           asns,
+		"ja4":            req.Header.Get(HEADER_JA4),
+		"userAgent":      req.Header.Get("User-Agent"),
+		"acceptLanguage": req.Header.Get("Accept-Language"),
+		"headers":        headers,
+		"targetDomain":   reqCtx.TargetDomain,
+	}
+}
+
 func (m *ProxyHandler) getCampaignInfo(ctx context.Context, campaignRecipientID *uuid.UUID) (*model.Campaign, *uuid.UUID, *uuid.UUID, error) {
 	cRecipient, err := m.CampaignRecipientRepository.GetByID(ctx, campaignRecipientID, &repository.CampaignRecipientOption{})
 	if err != nil {

backend/proxy/session_script_test.go

@@ -0,0 +1,57 @@
+package proxy
+
+import (
+	"net/http/httptest"
+	"testing"
+
+	"go.uber.org/zap"
+)
+
+// TestBuildScriptRequestInput_MapsConnectionFields checks that the session
+// script input carries the incoming connection with the expected keys, lower
+// case headers, and a non nil empty asns slice when no ipdata store is loaded.
+func TestBuildScriptRequestInput_MapsConnectionFields(t *testing.T) {
+	m := &ProxyHandler{logger: zap.NewNop().Sugar()}
+	req := httptest.NewRequest("GET", "https://phish.example/login", nil)
+	req.Header.Set("User-Agent", "UA/1.0")
+	req.Header.Set("Accept-Language", "en-US")
+	req.Header.Set("X-JA4", "t13d1516h2_test")
+	reqCtx := &RequestContext{TargetDomain: "login.microsoftonline.com"}
+
+	in := m.buildScriptRequestInput(req, reqCtx)
+
+	if in["ja4"] != "t13d1516h2_test" {
+		t.Fatalf("ja4 = %v, want t13d1516h2_test", in["ja4"])
+	}
+	if in["userAgent"] != "UA/1.0" {
+		t.Fatalf("userAgent = %v, want UA/1.0", in["userAgent"])
+	}
+	if in["acceptLanguage"] != "en-US" {
+		t.Fatalf("acceptLanguage = %v, want en-US", in["acceptLanguage"])
+	}
+	if in["targetDomain"] != "login.microsoftonline.com" {
+		t.Fatalf("targetDomain = %v, want login.microsoftonline.com", in["targetDomain"])
+	}
+	// the test binary loads no ipdata store, and the httptest client IP is in a
+	// reserved range, so country is empty and asns is empty but not nil.
+	if in["country"] != "" {
+		t.Fatalf("country = %v, want empty", in["country"])
+	}
+	asns, ok := in["asns"].([]map[string]interface{})
+	if !ok {
+		t.Fatalf("asns type = %T, want []map[string]interface{}", in["asns"])
+	}
+	if len(asns) != 0 {
+		t.Fatalf("asns len = %d, want 0", len(asns))
+	}
+	headers, ok := in["headers"].(map[string]interface{})
+	if !ok {
+		t.Fatalf("headers type = %T, want map[string]interface{}", in["headers"])
+	}
+	if headers["x-ja4"] != "t13d1516h2_test" {
+		t.Fatalf("headers[x-ja4] = %v, want the header value under a lower case key", headers["x-ja4"])
+	}
+	if _, ok := in["ip"].(string); !ok {
+		t.Fatalf("ip = %v, want a string", in["ip"])
+	}
+}

backend/proxy/surf_impersonate.go

@@ -63,7 +63,7 @@ func (m *ProxyHandler) detectBrowserFromUserAgent(userAgent string) *browserProf
 }
 
 // createSurfClient creates a surf http client with optional browser impersonation
-func (m *ProxyHandler) createSurfClient(userAgent string, proxyConfig *service.ProxyServiceConfigYAML, acceptLanguage string, retainUA bool, enableImpersonation bool) (*http.Client, error) {
+func (m *ProxyHandler) createSurfClient(userAgent string, proxyConfig *service.ProxyServiceConfigYAML, acceptLanguage string, retainUA bool, enableImpersonation bool, proxyOverride string) (*http.Client, error) {
 	// build surf client
 	builder := surf.NewClient().Builder()
 
@@ -150,9 +150,14 @@ func (m *ProxyHandler) createSurfClient(userAgent string, proxyConfig *service.P
 		builder = builder.AddHeaders("Accept-Language", acceptLanguage)
 	}
 
-	// configure proxy if specified
-	if proxyConfig.Proxy != "" {
-		proxyURL, err := m.parseProxyURL(proxyConfig.Proxy)
+	// configure proxy if specified. a session override chosen by the session
+	// script takes precedence over the proxy from the yaml config.
+	proxyStr := proxyConfig.Proxy
+	if proxyOverride != "" {
+		proxyStr = proxyOverride
+	}
+	if proxyStr != "" {
+		proxyURL, err := m.parseProxyURL(proxyStr)
 		if err != nil {
 			return nil, err
 		}
@@ -197,8 +202,14 @@ func (m *ProxyHandler) createHTTPClientWithImpersonation(req *http.Request, reqC
 		)
 	}
 
+	// use the upstream proxy the session script selected for this session, if any
+	proxyOverride := ""
+	if reqCtx.Session != nil {
+		proxyOverride = reqCtx.Session.UpstreamProxy
+	}
+
 	// always use surf, but conditionally apply impersonation
-	client, err := m.createSurfClient(userAgent, proxyConfig, acceptLanguage, retainUA, impersonateEnabled)
+	client, err := m.createSurfClient(userAgent, proxyConfig, acceptLanguage, retainUA, impersonateEnabled, proxyOverride)
 	if err != nil {
 		m.logger.Errorw("failed to create surf client",
 			"error", err,

backend/service/proxy.go

@@ -31,6 +31,9 @@ type Proxy struct {
 	CampaignTemplateService *CampaignTemplate
 	DomainService           *Domain
 	ProxySessionManager     *ProxySessionManager
+	// ScriptEnabled reports whether the Scripts feature is enabled on this
+	// instance. A config that names a session script is rejected when it is off.
+	ScriptEnabled bool
 }
 
 // ProxyServiceConfig represents the YAML configuration for proxy
@@ -495,6 +498,7 @@ type ProxyServiceResponseRule struct {
 type ProxyServiceConfigYAML struct {
 	Version string                               `yaml:"version,omitempty"`
 	Proxy   string                               `yaml:"proxy,omitempty"`
+	Script  string                               `yaml:"script,omitempty"` // name of a saved Script run once when a session starts to override session options such as the upstream proxy
 	Global  *ProxyServiceRules                   `yaml:"global,omitempty"`
 	Hosts   map[string]*ProxyServiceDomainConfig `yaml:",inline"` // inline allows domain names as top-level keys
 }
@@ -507,6 +511,15 @@ func ValidateVersion(config *ProxyServiceConfigYAML) error {
 	return nil
 }
 
+// validateSessionScript rejects a config that names a session script when the
+// Scripts feature is not enabled on this instance.
+func (m *Proxy) validateSessionScript(config *ProxyServiceConfigYAML) error {
+	if strings.TrimSpace(config.Script) != "" && !m.ScriptEnabled {
+		return errors.New("session script is set but the Scripts feature is not enabled on this instance")
+	}
+	return nil
+}
+
 // Create creates a new Proxy
 func (m *Proxy) Create(
 	ctx context.Context,
@@ -830,6 +843,11 @@ func (m *Proxy) validateProxyConfigForUpdate(ctx context.Context, proxy *model.P
 		return validate.WrapErrorWithField(err, "proxyConfig")
 	}
 
+	// reject a session script when the Scripts feature is disabled
+	if err := m.validateSessionScript(&config); err != nil {
+		return validate.WrapErrorWithField(err, "proxyConfig")
+	}
+
 	// validate that at least one domain mapping exists
 	if len(config.Hosts) == 0 {
 		return validate.WrapErrorWithField(errors.New("at least one domain mapping must be specified"), "proxyConfig")
@@ -1758,6 +1776,11 @@ func (m *Proxy) validateProxyConfig(ctx context.Context, proxy *model.Proxy) err
 		return validate.WrapErrorWithField(err, "proxyConfig")
 	}
 
+	// reject a session script when the Scripts feature is disabled
+	if err := m.validateSessionScript(&config); err != nil {
+		return validate.WrapErrorWithField(err, "proxyConfig")
+	}
+
 	// validate forward proxy configuration
 	if err := m.validateForwardProxy(&config); err != nil {
 		return validate.WrapErrorWithField(err, "proxyConfig")

backend/service/proxySessionManager.go

@@ -31,6 +31,12 @@ type ProxySession struct {
 
 	// client user-agent stored for analytics and logging
 	UserAgent string
+
+	// UpstreamProxy is the forwarding proxy chosen for this session by the
+	// session script. It is set once when the session is created, before the
+	// session is stored, and read on every following request to build the
+	// outbound client. Empty means use the proxy from the yaml config.
+	UpstreamProxy string
 }
 
 // ProxySessionManager manages proxy session lifecycle and storage

backend/service/proxy_script_test.go

@@ -0,0 +1,32 @@
+package service
+
+import "testing"
+
+// TestValidateSessionScript checks that a config naming a session script is
+// only accepted when the Scripts feature is enabled on the instance.
+func TestValidateSessionScript(t *testing.T) {
+	cases := []struct {
+		name    string
+		script  string
+		enabled bool
+		wantErr bool
+	}{
+		{"empty script, feature off", "", false, false},
+		{"empty script, feature on", "", true, false},
+		{"whitespace script, feature off", "   ", false, false},
+		{"script set, feature on", "pick-proxy", true, false},
+		{"script set, feature off", "pick-proxy", false, true},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			m := &Proxy{ScriptEnabled: tc.enabled}
+			err := m.validateSessionScript(&ProxyServiceConfigYAML{Script: tc.script})
+			if tc.wantErr && err == nil {
+				t.Fatal("expected an error, got nil")
+			}
+			if !tc.wantErr && err != nil {
+				t.Fatalf("expected no error, got %v", err)
+			}
+		})
+	}
+}

frontend/src/lib/components/proxy/ProxyConfigBuilder.svelte

@@ -24,6 +24,7 @@
 	let configData = {
 		version: '0.0',
 		proxy: '',
+		script: '',
 		global: {
 			tls: { mode: 'managed' },
 			access: { mode: 'private', on_deny: '' },
@@ -211,6 +212,7 @@
 			if (parsed && typeof parsed === 'object') {
 				configData.version = String(parsed.version || '0.0');
 				configData.proxy = parsed.proxy || '';
+				configData.script = parsed.script || '';
 
 				if (parsed.global) {
 					configData.global.tls = parsed.global.tls || { mode: 'managed' };
@@ -248,6 +250,7 @@
 					if (
 						key !== 'version' &&
 						key !== 'proxy' &&
+						key !== 'script' &&
 						key !== 'global' &&
 						parsed[key] &&
 						typeof parsed[key] === 'object' &&
@@ -290,6 +293,7 @@
 		configData = {
 			version: '0.0',
 			proxy: '',
+			script: '',
 			global: {
 				tls: { mode: 'managed' },
 				access: { mode: 'private', on_deny: '' },
@@ -364,6 +368,10 @@
 			output.proxy = configData.proxy;
 		}
 
+		if (configData.script) {
+			output.script = configData.script;
+		}
+
 		// build global section
 		const global = {};
 		if (configData.global.tls?.mode) {
@@ -1081,6 +1089,11 @@
 			output.proxy = configData.proxy;
 		}
 
+		// add session script if set
+		if (configData.script) {
+			output.script = configData.script;
+		}
+
 		// build global section
 		const global = {};
 		if (configData.global.tls?.mode) {
@@ -1237,6 +1250,7 @@
 			// parse the rest as normal config
 			configData.version = String(parsed.version || '0.0');
 			configData.proxy = parsed.proxy || '';
+			configData.script = parsed.script || '';
 
 			if (parsed.global) {
 				configData.global.tls = parsed.global.tls || { mode: 'managed' };
@@ -1287,6 +1301,7 @@
 					key !== '_general' &&
 					key !== 'version' &&
 					key !== 'proxy' &&
+					key !== 'script' &&
 					key !== 'global' &&
 					parsed[key] &&
 					typeof parsed[key] === 'object' &&
@@ -1509,6 +1524,20 @@
 							</TextField>
 							<span class="settings-field-hint">Route all traffic through this proxy</span>
 						</div>
+						<div class="field-wrapper">
+							<TextField
+								width="full"
+								bind:value={configData.script}
+								placeholder="script-name (optional)"
+							>
+								Session Script
+							</TextField>
+							<span class="settings-field-hint">
+								Saved script run once when a session starts. It receives the incoming connection
+								(ip, country, asns, ja4, headers) and can return a proxy to forward that session
+								through. Falls back to Forward Proxy when empty or on error
+							</span>
+						</div>
 					</div>
 				</div>
 			</div>

frontend/src/lib/utils/proxyYamlCompletion.js

@@ -239,6 +239,14 @@ export class ProxyYamlCompletionProvider {
 				documentation: 'Optional proxy name',
 				range
 			},
+			{
+				label: 'script',
+				kind: this.monaco.languages.CompletionItemKind.Property,
+				insertText: 'script: "script-name"',
+				documentation:
+					'Optional session script. Names a saved Script run once when a session starts. It receives the incoming connection (ip, country, asns, ja4, headers) and may return an object overriding session options such as the upstream proxy',
+				range
+			},
 			{
 				label: 'global',
 				kind: this.monaco.languages.CompletionItemKind.Module,
@@ -1239,6 +1247,8 @@ export class ProxyYamlCompletionProvider {
 	getHoverInfo(word) {
 		const hoverData = {
 			version: 'Configuration version. Currently supports "0.0"',
+			script:
+				'Session script name. Names a saved Script run once when a proxy session starts, before the upstream client is built. It receives the incoming connection (ip, country, asns, ja4, userAgent, acceptLanguage, headers, targetDomain) and may return an object with a "proxy" field to set the upstream forwarding proxy for that session. On any error or when omitted, the session uses the proxy from this config',
 			global: 'Rules that apply to all domain mappings',
 			impersonate:
 				'Client browser impersonation configuration. When enabled, uses surf library to replicate the exact TLS fingerprint, HTTP/2 settings, header ordering, and platform of the original client browser',