MCPX MCP Gateway

import yaml from dataclasses import dataclass, field from typing import Any, Optional from utils.consts import * from toolkit_testing.integration_tests.docker import write_file, read_file @dataclass class DomainLists: allowed_domains: list[str] = field(default_factory=lambda: []) blocked_domains: list[str] = field(default_factory=lambda: []) @dataclass class GlobalPolicy: remedies: list[Any] = field(default_factory=lambda: []) diagnosis: list[Any] = field(default_factory=lambda: []) @dataclass class EndpointPolicy: method: str url: str remedies: list[Any] = field(default_factory=lambda: []) diagnosis: list[Any] = field(default_factory=lambda: []) @dataclass class Header: name: Optional[str] value: Optional[str] @dataclass class Body: name: Optional[str] value: Optional[str] @dataclass class Token: header: Header @dataclass class OAuth: tokens: list[Body] = field(default_factory=lambda: []) @dataclass class BasicAuth: username: Optional[str] password: Optional[str] @dataclass class APIKey: tokens: list[Header] = field(default_factory=lambda: []) @dataclass class Authentication: o_auth: Optional[OAuth] = field(default=None) api_key: Optional[APIKey] = field(default=None) basic: Optional[BasicAuth] = field(default=None) @dataclass class Account: tokens: list[Token] = field(default_factory=lambda: []) authentication: Optional[Authentication] = field(default=None) @dataclass class PoliciesRequests: # should be `global`, however it is a reserved word in Python domain_lists: DomainLists = field(default_factory=lambda: DomainLists()) globals: GlobalPolicy = field(default_factory=lambda: GlobalPolicy()) endpoints: list[EndpointPolicy] = field(default_factory=lambda: []) accounts: dict[str, Account] = field(default_factory=lambda: {}) exporters: dict[str, Any] = field(default_factory=lambda: {}) def build_yaml(self) -> dict[str, Any]: endpoints_by_endpoint: dict[tuple[str, str], list[EndpointPolicy]] = {} for endpoint in self.endpoints: if not (endpoint.method, endpoint.url) in endpoints_by_endpoint: endpoints_by_endpoint[(endpoint.method, endpoint.url)] = [] endpoints_by_endpoint[(endpoint.method, endpoint.url)].append(endpoint) def merge_endpoint_policies( method: str, url: str, policies: list[EndpointPolicy] ) -> EndpointPolicy: merged_policy = EndpointPolicy( method=method, url=url, remedies=[], diagnosis=[] ) for policy in policies: merged_policy.remedies.extend(policy.remedies) merged_policy.diagnosis.extend(policy.diagnosis) return merged_policy merged_endpoints = [ merge_endpoint_policies(k[0], k[1], v) for (k, v) in endpoints_by_endpoint.items() ] return { "allowed_domains": self.domain_lists.allowed_domains, "blocked_domains": self.domain_lists.blocked_domains, "global": self.globals, "endpoints": merged_endpoints, "accounts": self.accounts, "exporters": self.exporters, } async def read_actual_policies_file() -> dict[str, Any]: return await _read_policies_file( policies_filename=PoliciesFilename.ACTUAL_POLICIES_FILENAME ) async def read_initial_policies_file() -> dict[str, Any]: return await _read_policies_file( policies_filename=PoliciesFilename.INITIAL_POLICIES_FILENAME ) async def _read_policies_file( policies_filename: PoliciesFilename, container_name: str = LUNAR_PROXY_SERVICE_NAME ) -> dict[str, Any]: existing_policies_file = await read_file( container_name=container_name, directory_path=POLICIES_DIRECTORY, file=policies_filename.value, ) return yaml.safe_load(existing_policies_file) async def write_actual_policies_file( policies_yaml: dict[str, Any], container_name: str = LUNAR_PROXY_SERVICE_NAME ): return await _write_policies_file( policies_filename=PoliciesFilename.ACTUAL_POLICIES_FILENAME, policies_yaml_str=yaml.dump(policies_yaml, default_flow_style=False), container_name=container_name, ) async def write_initial_policies_file( policies_yaml: str, container_name: str = LUNAR_PROXY_SERVICE_NAME ): return await _write_policies_file( policies_filename=PoliciesFilename.INITIAL_POLICIES_FILENAME, policies_yaml_str=policies_yaml, container_name=container_name, ) async def _write_policies_file( policies_filename: PoliciesFilename, policies_yaml_str: str, container_name: str = LUNAR_PROXY_SERVICE_NAME, ): await write_file( container_name=container_name, directory_path=POLICIES_DIRECTORY, file=policies_filename.value, content=policies_yaml_str, )