feat: introduce ProxyProtocol enum and update related classes to utilize it

Author: M03ED

app/core/abstract_core.py

@@ -1,5 +1,7 @@
 from abc import ABC, abstractmethod
 
+from app.models.protocol import ProxyProtocol
+
 
 class AbstractCore(ABC):
     @abstractmethod
@@ -35,3 +37,8 @@ def from_json(cls, data: dict) -> "AbstractCore":
     @abstractmethod
     def inbounds(self) -> list[str]:
         raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def protocols(self) -> frozenset[ProxyProtocol]:
+        raise NotImplementedError

app/core/wireguard.py

@@ -9,8 +9,11 @@
 import commentjson
 
 from app.models.core import CoreType
+from app.models.protocol import ProxyProtocol
 from app.utils.crypto import get_wireguard_public_key, validate_wireguard_key
 
+_WIREGUARD_PROTOCOLS = frozenset((ProxyProtocol.wireguard,))
+
 
 class WireGuardConfig(dict):
     def __init__(
@@ -115,6 +118,10 @@ def inbounds_by_tag(self) -> dict:
     def inbounds(self) -> list[str]:
         return self._inbounds
 
+    @property
+    def protocols(self) -> frozenset[ProxyProtocol]:
+        return _WIREGUARD_PROTOCOLS
+
     def to_json(self) -> dict:
         return {
             "type": self.type,

app/core/xray.py

@@ -9,9 +9,18 @@
 import commentjson
 
 from app.models.core import CoreType
+from app.models.protocol import ProxyProtocol
 from app.utils.crypto import get_cert_SANs, get_x25519_public_key
 
 
+def _protocols_from_inbounds_by_tag(inbounds_by_tag: dict[str, dict]) -> frozenset[ProxyProtocol]:
+    return frozenset(
+        protocol
+        for inbound in inbounds_by_tag.values()
+        if (protocol := ProxyProtocol.from_value(inbound["protocol"])) is not None
+    )
+
+
 class XRayConfig(dict):
     def __init__(
         self,
@@ -45,6 +54,7 @@ def __init__(
         self._inbounds = []
         self._inbounds_by_tag = {}
         self._fallbacks_inbound = []
+        self._protocols: frozenset[ProxyProtocol] = frozenset()
 
         # Registery pattern for network handlers, making it easy to add support for new network types in the future
         self.network_handlers = {
@@ -360,6 +370,7 @@ def _resolve_inbounds(self):
         """Resolve all inbounds and their settings."""
         for inbound in self["inbounds"]:
             self._read_inbound(inbound)
+        self._protocols = _protocols_from_inbounds_by_tag(self._inbounds_by_tag)
 
     def _read_inbound(self, inbound: dict):
         """Read an inbound and its settings."""
@@ -471,6 +482,10 @@ def inbounds(self) -> list[str]:
         """Get inbounds by tag."""
         return self._inbounds
 
+    @property
+    def protocols(self) -> frozenset[ProxyProtocol]:
+        return self._protocols
+
     @property
     def type(self) -> str:
         return self._type
@@ -503,6 +518,7 @@ def from_json(cls, data: dict) -> "XRayConfig":
             instance._inbounds = data["inbounds"]
         if "inbounds_by_tag" in data:
             instance._inbounds_by_tag = data["inbounds_by_tag"]
+        instance._protocols = _protocols_from_inbounds_by_tag(instance._inbounds_by_tag)
         return instance
 
     def copy(self):

app/models/protocol.py

@@ -0,0 +1,20 @@
+from enum import IntEnum
+
+
+class ProxyProtocol(IntEnum):
+    vmess = 1
+    vless = 2
+    trojan = 3
+    shadowsocks = 4
+    wireguard = 5
+    hysteria = 6
+
+    @classmethod
+    def from_value(cls, value: str) -> "ProxyProtocol" | None:
+        try:
+            return _PROXY_PROTOCOL_BY_NAME[value]
+        except KeyError:
+            return None
+
+
+_PROXY_PROTOCOL_BY_NAME = {protocol.name: protocol for protocol in ProxyProtocol}

app/node/user.py

@@ -1,13 +1,12 @@
-import inspect
-
 from PasarGuardNodeBridge import create_proxy, create_user
 from PasarGuardNodeBridge.common.service_pb2 import User as ProtoUser
 from sqlalchemy import and_, func, select
 
 from app.db import AsyncSession
 from app.db.models import Group, ProxyInbound, User, UserStatus, inbounds_groups_association, users_groups_association
+from app.models.protocol import ProxyProtocol
 
-_CREATE_PROXY_PARAMS = set(inspect.signature(create_proxy).parameters)
+_ALL_PROXY_PROTOCOLS = frozenset(ProxyProtocol)
 
 
 def _inbounds_from_loaded_groups(user: User) -> list[str] | None:
@@ -30,7 +29,7 @@ def _inbounds_from_loaded_groups(user: User) -> list[str] | None:
     return list(tags)
 
 
-async def serialize_user(user: User) -> ProtoUser:
+async def serialize_user(user: User, allowed_protocols: frozenset[ProxyProtocol] | None = None) -> ProtoUser:
     user_settings = user.proxy_settings
     inbounds = None
     status = user.__dict__.get("status")
@@ -42,38 +41,52 @@ async def serialize_user(user: User) -> ProtoUser:
         if inbounds is None:
             inbounds = await user.inbounds()
 
-    return _serialize_user_for_node(user.id, user.username, user_settings, inbounds)
-
-
-def _serialize_user_for_node(id: int, username: str, user_settings: dict, inbounds: list[str] = None) -> ProtoUser:
-    vmess_settings = user_settings.get("vmess", {})
-    vless_settings = user_settings.get("vless", {})
-    if vless_settings.get("flow") == "xtls-rprx-vision-udp443":
-        vless_settings["flow"] = "xtls-rprx-vision"
-    trojan_settings = user_settings.get("trojan", {})
-    shadowsocks_settings = user_settings.get("shadowsocks", {})
-    wireguard_settings = user_settings.get("wireguard", {})
-    hysteria_settings = user_settings.get("hysteria", {})
-    proxy_kwargs = {
-        "vmess_id": vmess_settings.get("id"),
-        "vless_id": vless_settings.get("id"),
-        "vless_flow": vless_settings.get("flow"),
-        "trojan_password": trojan_settings.get("password"),
-        "shadowsocks_password": shadowsocks_settings.get("password"),
-        "shadowsocks_method": shadowsocks_settings.get("method"),
-        "wireguard_public_key": wireguard_settings.get("public_key"),
-        "wireguard_peer_ips": wireguard_settings.get("peer_ips") or [],
-        "hysteria_auth": hysteria_settings.get("auth"),
-    }
+    return _serialize_user_for_node(user.id, user.username, user_settings, inbounds, allowed_protocols)
+
+
+def _serialize_user_for_node(
+    id: int,
+    username: str,
+    user_settings: dict,
+    inbounds: list[str] = None,
+    allowed_protocols: frozenset[ProxyProtocol] | None = None,
+) -> ProtoUser:
+    allowed_protocols = allowed_protocols or _ALL_PROXY_PROTOCOLS
+
+    proxy_kwargs = {}
+    if ProxyProtocol.vmess in allowed_protocols:
+        proxy_kwargs["vmess_id"] = user_settings.get("vmess", {}).get("id")
+    if ProxyProtocol.vless in allowed_protocols:
+        vless_settings = dict(user_settings.get("vless", {}))
+        if vless_settings.get("flow") == "xtls-rprx-vision-udp443":
+            vless_settings["flow"] = "xtls-rprx-vision"
+        proxy_kwargs["vless_id"] = vless_settings.get("id")
+        proxy_kwargs["vless_flow"] = vless_settings.get("flow")
+    if ProxyProtocol.trojan in allowed_protocols:
+        proxy_kwargs["trojan_password"] = user_settings.get("trojan", {}).get("password")
+    if ProxyProtocol.shadowsocks in allowed_protocols:
+        shadowsocks_settings = user_settings.get("shadowsocks", {})
+        proxy_kwargs["shadowsocks_password"] = shadowsocks_settings.get("password")
+        proxy_kwargs["shadowsocks_method"] = shadowsocks_settings.get("method")
+    if ProxyProtocol.wireguard in allowed_protocols:
+        wireguard_settings = user_settings.get("wireguard", {})
+        proxy_kwargs["wireguard_public_key"] = wireguard_settings.get("public_key")
+        proxy_kwargs["wireguard_peer_ips"] = wireguard_settings.get("peer_ips") or []
+    if ProxyProtocol.hysteria in allowed_protocols:
+        proxy_kwargs["hysteria_auth"] = user_settings.get("hysteria", {}).get("auth")
 
     return create_user(
         f"{id}.{username}",
-        create_proxy(**{key: value for key, value in proxy_kwargs.items() if key in _CREATE_PROXY_PARAMS}),
+        create_proxy(**proxy_kwargs),
         inbounds,
     )
 
 
-async def core_users(db: AsyncSession, inbound_tags: list[str] | set[str] | None = None):
+async def core_users(
+    db: AsyncSession,
+    inbound_tags: list[str] | set[str] | None = None,
+    allowed_protocols: frozenset[ProxyProtocol] | None = None,
+):
     dialect = db.bind.dialect.name
     inbound_tags = list(dict.fromkeys(inbound_tags or []))
 
@@ -117,11 +130,21 @@ async def core_users(db: AsyncSession, inbound_tags: list[str] | set[str] | None
     for row in results:
         inbound_tags = row.inbound_tags.split(",") if row.inbound_tags else []
         if inbound_tags:
-            bridge_users.append(_serialize_user_for_node(row.id, row.username, row.proxy_settings, inbound_tags))
+            bridge_users.append(
+                _serialize_user_for_node(
+                    row.id,
+                    row.username,
+                    row.proxy_settings,
+                    inbound_tags,
+                    allowed_protocols,
+                )
+            )
     return bridge_users
 
 
-async def serialize_users_for_node(users: list[User]) -> list[ProtoUser]:
+async def serialize_users_for_node(
+    users: list[User], allowed_protocols: frozenset[ProxyProtocol] | None = None
+) -> list[ProtoUser]:
     bridge_users: list = []
 
     for user in users:
@@ -133,6 +156,8 @@ async def serialize_users_for_node(users: list[User]) -> list[ProtoUser]:
             else:
                 inbounds_list = loaded_inbounds
 
-        bridge_users.append(_serialize_user_for_node(user.id, user.username, user.proxy_settings, inbounds_list))
+        bridge_users.append(
+            _serialize_user_for_node(user.id, user.username, user.proxy_settings, inbounds_list, allowed_protocols)
+        )
 
     return bridge_users

app/operation/node.py

@@ -209,7 +209,11 @@ async def _get_core_users_map(
                 users_by_core[core_id] = []
                 continue
 
-            users_by_core[core_id] = await core_users(db=db, inbound_tags=core.inbounds)
+            users_by_core[core_id] = await core_users(
+                db=db,
+                inbound_tags=core.inbounds,
+                allowed_protocols=core.protocols,
+            )
 
         return cores_by_id, users_by_core