From 0641864478066df0a918f87103744f06ad0c1217 Mon Sep 17 00:00:00 2001
From: Abhinav Singh <126065+abhinavsingh@users.noreply.github.com>
Date: Sun, 13 Oct 2024 15:56:16 +0530
Subject: [PATCH] `GroutClientBasePlugin` and example `GroutClientPlugin`
 (#1488)

* `GroutClientBasePlugin`

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Silence `S113` false positive

* Remove example url

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .flake8                      |  1 +
 proxy/common/constants.py    | 11 ++++++-----
 proxy/plugin/grout_client.py | 33 +++++++++++++++++++++++++++++++++
 proxy/proxy.py               | 28 ++++++++++++++++++++++++++++
 4 files changed, 68 insertions(+), 5 deletions(-)
 create mode 100644 proxy/plugin/grout_client.py

diff --git a/.flake8 b/.flake8
index 39cf6840c1..cfc0b078eb 100644
--- a/.flake8
+++ b/.flake8
@@ -82,6 +82,7 @@ extend-ignore =
   S101    # FIXME: assertions are thrown away in optimized mode, needs audit
   S104    # FIXME: bind-all interface listen
   S105    # FIXME: hardcoded password?
+  S113    # FIXME: Call to httpx without timeout (false positive)
   S303    # FIXME: insecure hash func
   S311    # FIXME: `random` needs auditing
   S404    # FIXME: `subprocess` use needs auditing
diff --git a/proxy/common/constants.py b/proxy/common/constants.py
index 90bcaf9fec..f1385390b3 100644
--- a/proxy/common/constants.py
+++ b/proxy/common/constants.py
@@ -180,11 +180,12 @@ def _env_threadless_compliant() -> bool:
 
 # Cor plugins enabled by default or via flags
 DEFAULT_ABC_PLUGINS = [
-    'HttpProtocolHandlerPlugin',
-    'HttpProxyBasePlugin',
-    'HttpWebServerBasePlugin',
-    'WebSocketTransportBasePlugin',
-    'ReverseProxyBasePlugin',
+    "HttpProtocolHandlerPlugin",
+    "HttpProxyBasePlugin",
+    "HttpWebServerBasePlugin",
+    "WebSocketTransportBasePlugin",
+    "ReverseProxyBasePlugin",
+    "GroutClientBasePlugin",
 ]
 PLUGIN_DASHBOARD = 'proxy.dashboard.ProxyDashboard'
 PLUGIN_HTTP_PROXY = 'proxy.http.proxy.HttpProxyPlugin'
diff --git a/proxy/plugin/grout_client.py b/proxy/plugin/grout_client.py
new file mode 100644
index 0000000000..109395123e
--- /dev/null
+++ b/proxy/plugin/grout_client.py
@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+"""
+    proxy.py
+    ~~~~~~~~
+    ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
+    Network monitoring, controls & Application development, testing, debugging.
+
+    :copyright: (c) 2013-present by Abhinav Singh and contributors.
+    :license: BSD, see LICENSE for more details.
+"""
+
+
+from proxy.proxy import GroutClientBasePlugin
+from proxy.common.types import HostPort
+from proxy.http.parser.parser import HttpParser
+
+
+class GroutClientPlugin(GroutClientBasePlugin):
+
+    def resolve_route(
+        self,
+        route: str,
+        request: HttpParser,
+        origin: HostPort,
+        server: HostPort,
+    ) -> str:
+        print(request, origin, server, '->', route)
+        print(request.header(b'host'), request.path)
+        # Send to localhost:7001 irrespective of the
+        # original "route" value provided to the grout client
+        # OR any custom host:upstream mapping provided through the
+        # --tunnel-route flags.
+        return 'http://localhost:7001'
diff --git a/proxy/proxy.py b/proxy/proxy.py
index 3dbbc9377d..ce5509696b 100644
--- a/proxy/proxy.py
+++ b/proxy/proxy.py
@@ -20,6 +20,7 @@
 import logging
 import argparse
 import threading
+from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, List, Type, Tuple, Optional, cast
 
 from .core.ssh import SshTunnelListener, SshHttpProtocolHandler
@@ -28,6 +29,7 @@
 from .http.codes import httpStatusCodes
 from .common.flag import FlagParser, flags
 from .http.client import client
+from .common.types import HostPort
 from .common.utils import bytes_
 from .core.work.fd import RemoteFdExecutor
 from .http.methods import httpMethods
@@ -44,6 +46,7 @@
     DEFAULT_SSH_LISTENER_KLASS,
 )
 from .core.event.metrics import MetricsEventSubscriber
+from .http.parser.parser import HttpParser
 
 
 if TYPE_CHECKING:   # pragma: no cover
@@ -509,3 +512,28 @@ def _parse() -> Tuple[str, int]:
     assert env is not None
     print('\r' + ' ' * 70 + '\r', end='', flush=True)
     Plugins.from_bytes(env['m'].encode(), name='client').grout(env=env['e'])  # type: ignore[attr-defined]
+
+
+class GroutClientBasePlugin(ABC):
+    """Base class for dynamic grout client rules.
+
+    Implementation of this class must be stateless because a new instance is created
+    for every route decision making.
+    """
+
+    @abstractmethod
+    def resolve_route(
+        self,
+        route: str,
+        request: HttpParser,
+        origin: HostPort,
+        server: HostPort,
+    ) -> str:
+        """Returns a valid grout route string.
+
+        You MUST override this method.  For a simple pass through,
+        simply return the "route" argument value itself.  You can also
+        return a dynamic value based upon "request" and "origin" information.
+        E.g. sending to different upstream services based upon request Host header.
+        """
+        raise NotImplementedError()