fmtr.tools 1.1.1__py3-none-any.whl → 1.3.81__py3-none-any.whl

fmtr/tools/dns_tools/dm.py

@@ -0,0 +1,257 @@
+import dns
+import httpx
+from dataclasses import dataclass, field
+from dns import rcode as dnspython_rcode, reversename as dnspython_reversename
+from dns.message import Message, QueryMessage
+from dns.rrset import RRset
+from functools import cached_property
+from typing import Self, Optional, List
+
+from fmtr.tools.string_tools import join
+
+TTL_CODE_DEFAULTS = {
+    dnspython_rcode.NOERROR: 300,  # Successful query
+    dnspython_rcode.FORMERR: 60,  # Format error
+    dnspython_rcode.SERVFAIL: 10,  # Server failure
+    dnspython_rcode.NXDOMAIN: 60 * 60,  # Non-existent domain
+    dnspython_rcode.NOTIMP: 60,  # Not implemented
+    dnspython_rcode.REFUSED: 60,  # Refused
+    dnspython_rcode.YXDOMAIN: 600,  # Name exists when it should not
+    dnspython_rcode.YXRRSET: 600,  # RR Set exists when it should not
+    dnspython_rcode.NXRRSET: 300,  # RR Set that should exist does not
+    dnspython_rcode.NOTAUTH: 60,  # Not authorized
+    dnspython_rcode.NOTZONE: 60  # Name not contained in zone
+}
+
+@dataclass
+class BaseDNSData:
+    """
+
+    DNS response object.
+
+    """
+    wire: bytes
+
+    @cached_property
+    def message(self) -> Message:
+        return dns.message.from_wire(self.wire)
+
+    @classmethod
+    def from_message(cls, message: Message) -> Self:
+        return cls(message.to_wire())
+
+@dataclass
+class Response(BaseDNSData):
+    """
+
+    DNS response object.
+
+    """
+
+    http: Optional[httpx.Response] = None
+    blocked_by: Optional[str] = None
+
+    @classmethod
+    def from_http(cls, response: httpx.Response) -> Self:
+        """
+
+        Initialise from an HTTP response.
+
+        """
+        self = cls(response.content, http=response)
+        return self
+
+    @property
+    def answer(self) -> Optional[RRset]:
+        """
+
+        Get the latest answer, if one exists.
+
+        """
+        if not self.message.answer:
+            return None
+        return self.message.answer[-1]
+
+    @property
+    def rcode(self) -> dnspython_rcode.Rcode:
+        return self.message.rcode()
+
+    @property
+    def rcode_text(self) -> str:
+        return dnspython_rcode.to_text(self.rcode)
+
+    @property
+    def ttl(self) -> int:
+        """
+
+        Get median TTL from answers, falling back to authority, then to error-code defaults.
+
+        """
+        answers = self.message.answer or self.message.authority
+        if answers:
+            ttls = [answer.ttl for answer in answers]
+            ttl = min(ttls)
+            return ttl
+
+        ttl = TTL_CODE_DEFAULTS.get(self.rcode, dnspython_rcode.NXDOMAIN)
+        return ttl
+
+
+
+    def __str__(self):
+        """
+
+        Put answer and ID text in string representation.
+
+        """
+        answer = self.answer
+
+        if answer:
+            answer = join(answer.to_text().splitlines(), sep=', ')
+
+        string = join([answer, self.message.flags], sep=', ')
+        string = f'{self.__class__.__name__}({string})'
+        return string
+
+
+
+@dataclass
+class Request(BaseDNSData):
+    """
+
+    DNS request object.
+
+    """
+    wire: bytes
+
+    @cached_property
+    def question(self) -> RRset:
+        return self.message.question[0]
+
+    @cached_property
+    def is_valid(self):
+        return len(self.message.question) != 0
+
+    @cached_property
+    def type(self):
+        return self.question.rdtype
+
+    @cached_property
+    def type_text(self):
+        return dns.rdatatype.to_text(self.type)
+
+    @cached_property
+    def name(self):
+        return self.question.name
+
+    @cached_property
+    def name_text(self):
+        return self.name.to_text()
+
+    def get_response_template(self):
+        message = dns.message.make_response(self.message)
+        message.flags |= dns.flags.RA
+        return message
+
+    @cached_property
+    def blackhole(self) -> Response:
+        blackhole = self.get_response_template()
+        blackhole.set_rcode(dns.rcode.NXDOMAIN)
+        response = Response.from_message(blackhole)
+        return response
+
+
+@dataclass
+class Exchange:
+    """
+
+    Entire DNS exchange for a DNS Proxy: request -> upstream response -> response
+
+    """
+    ip: str
+    port: int
+
+    request: Request
+    response: Optional[Response] = None
+    answers_pre: List[RRset] = field(default_factory=list)
+    is_internal: bool = False
+    client_name: Optional[str] = None
+    is_complete: bool = False
+
+    @property
+    def addr(self):
+        return self.ip, self.port
+
+    @classmethod
+    def from_wire(cls, wire: bytes, **kwargs) -> Self:
+        request = Request(wire)
+        response = Response.from_message(request.get_response_template())
+        return cls(request=request, response=response, **kwargs)
+
+    @cached_property
+    def client(self):
+        return f'{self.ip}:{self.port}'
+
+    @property
+    def question_last(self) -> RRset:
+        """
+
+        Create an RRset surrogate representing the latest/current question.
+        This can be the original question - or a hybrid one if we've injected our own answers into the Exchange.
+        If there's a response, use its answers, else fall back to answers_pre, else to the original question.
+
+        """
+        answers = self.answers_pre
+        if self.response:
+            answers = self.response.message.answer or answers
+
+        if answers:
+            rrset = answers[-1]
+            rdtype = self.request.type
+            ttl = self.request.question.ttl
+            rdclass = self.request.question.rdclass
+            name = next(iter(rrset.items.keys())).to_text()
+            rrset_surrogate = dns.rrset.from_text(
+                name=name,
+                ttl=ttl,
+                rdtype=rdtype,
+                rdclass=rdclass,
+            )
+
+            return rrset_surrogate
+        else:
+            return self.request.question
+
+    @property
+    def query_last(self) -> QueryMessage:
+        """
+
+        Create a query (e.g. for use by upstream) based on the last question.
+
+        """
+
+        question_last = self.question_last
+        query = dns.message.make_query(qname=question_last.name, rdclass=question_last.rdclass, rdtype=question_last.rdtype, id=self.request.message.id)
+        return query
+
+    @property
+    def key(self):
+        """
+
+        Hashable key for caching
+
+        """
+        data = tuple(self.request.question.to_text().split())
+        return data
+
+    @cached_property
+    def reverse(self) -> Self:
+        """
+
+        Create an Exchange for a reverse lookup of this Exchange's client IP.
+
+        """
+        name = dnspython_reversename.from_address(self.ip)
+        query = dns.message.make_query(name, dns.rdatatype.PTR)
+        exchange = self.__class__.from_wire(query.to_wire(), ip=self.ip, port=self.port, is_internal=True)
+        return exchange

fmtr/tools/dns_tools/proxy.py

@@ -0,0 +1,66 @@
+from dataclasses import dataclass
+
+from fmtr.tools.dns_tools import server, client
+from fmtr.tools.dns_tools.dm import Exchange
+from fmtr.tools.logging_tools import logger
+
+
+@dataclass(kw_only=True, eq=False)
+class Proxy(server.Plain):
+    """
+
+    Base for a DNS Proxy server (plain server) TODO: Allow subclassing of any server type.
+
+    """
+
+    client: client.HTTP
+
+    def process_question(self, exchange: Exchange):
+        """
+
+        Modify exchange based on initial question.
+
+        """
+        return
+
+    def process_upstream(self, exchange: Exchange):
+        """
+
+        Modify exchange after upstream response.
+
+        """
+        return
+
+    def finalize(self, exchange: Exchange):
+        """
+
+        Finalize a still open exchange.
+
+        """
+        exchange.is_complete = True
+
+    async def resolve(self, exchange: Exchange) -> Exchange:
+        """
+
+        Resolve a request, processing each stage, initial question, upstream response etc.
+        Subclasses can override the relevant processing methods to implement custom behaviour.
+
+        """
+        with logger.span(f'Processing question...'):
+            self.process_question(exchange)
+        if exchange.is_complete:
+            return exchange
+
+        with logger.span(f'Making upstream request...'):
+            self.client.resolve(exchange)
+        if exchange.is_complete:
+            return exchange
+
+        with logger.span(f'Processing upstream response...'):
+            self.process_upstream(exchange)
+        if exchange.is_complete:
+            return exchange
+
+        self.finalize(exchange)
+
+        return exchange

fmtr/tools/dns_tools/server.py

@@ -0,0 +1,138 @@
+import asyncio
+from dataclasses import dataclass, field
+from datetime import timedelta
+from dns import rcode as dnspython_rcode
+from functools import cached_property
+from typing import Optional
+
+from fmtr.tools import caching_tools as caching
+from fmtr.tools.dns_tools.dm import Exchange
+from fmtr.tools.logging_tools import logger
+
+
+@dataclass(kw_only=True, eq=False)
+class Plain(asyncio.DatagramProtocol):
+    """
+
+    Async base class for a plain DNS server using asyncio DatagramProtocol.
+    """
+
+    host: str
+    port: int
+    transport: Optional[asyncio.DatagramTransport] = field(default=None, init=False)
+
+    @cached_property
+    def loop(self):
+        return asyncio.get_event_loop()
+
+
+    @cached_property
+    def cache(self):
+        """
+
+        Overridable cache.
+        """
+        cache = caching.TLRU(maxsize=1_024, ttu_static=timedelta(hours=1), desc='DNS Request')
+        return cache
+
+    def connection_made(self, transport: asyncio.DatagramTransport):
+        self.transport = transport
+        logger.info(f'Listening on {self.host}:{self.port}')
+
+    def datagram_received(self, data: bytes, addr):
+        ip, port = addr
+        exchange = Exchange.from_wire(data, ip=ip, port=port)
+        asyncio.create_task(self.handle(exchange))
+
+    async def start(self):
+        """
+
+        Start the async UDP server.
+        """
+
+        logger.info(f'Starting async DNS server on {self.host}:{self.port}...')
+        await self.loop.create_datagram_endpoint(
+            lambda: self,
+            local_addr=(self.host, self.port)
+        )
+        await asyncio.Future()  # Prevent exit by blocking forever
+
+    async def resolve(self, exchange: Exchange) -> Exchange:
+        """
+
+        To be defined in subclasses.
+
+        """
+        raise NotImplementedError
+
+    def check_cache(self, exchange: Exchange):
+        if exchange.key in self.cache:
+            logger.info(f'Request found in cache.')
+            exchange.response = self.cache[exchange.key]
+            exchange.response.message.id = exchange.request.message.id
+            exchange.is_complete = True
+
+    def get_span(self, exchange: Exchange):
+        """
+
+        Get handling span
+
+        """
+        request = exchange.request
+        span = logger.span(
+            f'Handling request {exchange.client_name=} {request.message.id=} {request.type_text} {request.name_text} {request.question=}...'
+        )
+        return span
+
+    def log_response(self, exchange: Exchange):
+        """
+
+        Log when resolution complete
+
+        """
+        request = exchange.request
+        response = exchange.response
+        logger.info(
+            f'Resolution complete {exchange.client_name=} {request.message.id=} {request.type_text} {request.name_text} {request.question=} {exchange.is_complete=} {response.rcode=} {response.rcode_text=} {response.answer=} {response.blocked_by=}...'
+        )
+
+    def log_dns_errors(self, exchange: Exchange):
+        """
+
+        Warn about any errors
+
+        """
+        if exchange.response.rcode != dnspython_rcode.NOERROR:
+            logger.warning(f'Error {exchange.response.rcode_text=}')
+
+    async def handle(self, exchange: Exchange):
+        """
+
+        Warn about any errors
+
+        """
+        if not exchange.request.is_valid:
+            raise ValueError(f'Only one question per request is supported. Got {len(exchange.request.question)} questions.')
+
+        if not exchange.is_internal:
+            await self.handle(exchange.reverse)
+            client_name = exchange.reverse.question_last.name.to_text()
+            if not exchange.reverse.response.answer:
+                logger.warning(f'Client name could not be resolved {client_name=}.')
+            exchange.client_name = client_name
+
+        with self.get_span(exchange):
+            with logger.span(f'Checking cache...'):
+                self.check_cache(exchange)
+
+            if not exchange.is_complete:
+                exchange = await self.resolve(exchange)
+                self.cache[exchange.key] = exchange.response
+
+            self.log_dns_errors(exchange)
+            self.log_response(exchange)
+
+        if exchange.is_internal:
+            return
+
+        self.transport.sendto(exchange.response.message.to_wire(), exchange.addr)

fmtr/tools/docker_tools/__init__.py

@@ -0,0 +1,6 @@
+from fmtr.tools.import_tools import MissingExtraMockModule
+
+try:
+    from python_on_whales import DockerClient
+except ModuleNotFoundError as exception:
+    DockerClient = MissingExtraMockModule('docker.client', exception)

fmtr/tools/entrypoints/cache_hfh.py

@@ -0,0 +1,3 @@
+def main():
+    from fmtr import tools
+    tools.hfh.main()

fmtr/tools/entrypoints/ep_test.py

@@ -0,0 +1,2 @@
+def main():
+    print('Ran test entrypoint.')

fmtr/tools/entrypoints/install_yamlscript.py

@@ -0,0 +1,8 @@
+def main():
+    """
+
+    Ensure YS binary is installed.
+
+    """
+    from fmtr.tools import yaml
+    yaml.install()

fmtr/tools/{console_script_tools.py → entrypoints/remote_debug_test.py}

@@ -1,9 +1,4 @@
-def cache_hfh():
-    from fmtr import tools
-    tools.hfh.main()
-
-
-def remote_debug_test():
+def main():
     """
 
     Test debugger connection

fmtr/tools/entrypoints/shell_debug.py

@@ -0,0 +1,8 @@
+def main():
+    """
+
+    Start debugger
+
+    """
+    from fmtr.tools import debug
+    debug.debug_shell()

fmtr/tools/environment_tools.py

@@ -28,7 +28,7 @@ def get_dict() -> Dict[str, str]:
     Return environment variables as a standard dictionary.
 
     """
-    environment_dict = dict(os.environ)
+    environment_dict = dict(sorted(dict(os.environ).items()))
     return environment_dict
 
 
@@ -76,4 +76,4 @@ get_date = get_getter(date.fromisoformat)
 get_datetime = get_getter(datetime.fromisoformat)
 get_path = get_getter(Path)
 
-IS_DEBUG = get(Constants.FMTR_LOG_LEVEL_KEY, None, converter=str.upper) == 'DEBUG'
+IS_DEV = get_bool(Constants.FMTR_DEV_KEY, default=False)

fmtr/tools/function_tools.py

@@ -1,4 +1,8 @@
-from typing import Tuple
+import functools
+import inspect
+from typing import Tuple, Callable, Self
+
+from fmtr.tools import context_tools
 
 
 def combine_args_kwargs(args: dict=None, kwargs: dict=None) -> dict:
@@ -31,3 +35,75 @@ def split_args_kwargs(args_kwargs: dict) -> Tuple[list, dict]:
     return args, kwargs
 
 
+class MethodDecorator:
+    """
+
+    Bound method decorator with overridable start/stop and context manager
+
+    """
+
+    CONTEXT_KEY = 'context'
+
+    def __init__(self):
+        """
+
+        Initialise the decorator itself with any arguments
+
+        """
+        self.func = None
+
+    def __call__(self, func: Callable) -> Self:
+        """
+
+        Add the (unbound) method.
+
+        """
+        self.func = func
+        functools.update_wrapper(self, func)
+        return self
+
+    def __get__(self, instance, owner):
+        """
+
+        Wrap bound method at runtime, call start/stop within context.
+
+        """
+        if instance is None:  # Class method called.
+            return self.func
+
+        if inspect.iscoroutinefunction(self.func):
+            async def async_wrapper(*args, **kwargs):
+                with self.get_context(instance):
+                    self.start(instance, *args, **kwargs)
+                    result = await self.func(instance, *args, **kwargs)
+                    self.stop(instance, *args, **kwargs)
+                    return result
+
+            return async_wrapper
+
+        else:
+            def sync_wrapper(*args, **kwargs):
+                with self.get_context(instance):
+                    self.start(instance, *args, **kwargs)
+                    result = self.func(instance, *args, **kwargs)
+                    self.stop(instance, *args, **kwargs)
+                    return result
+
+            return sync_wrapper
+
+    def get_context(self, instance):
+        """
+
+        If the instance has a context attribute, use that - otherwise use a null context.
+
+        """
+        context = getattr(instance, self.CONTEXT_KEY, None)
+        if context:
+            return context
+        return context_tools.null()
+
+    def start(self, instance, *args, **kwargs):
+        pass
+
+    def stop(self, instance, *args, **kwargs):
+        pass

fmtr/tools/google_api_tools.py

@@ -1,3 +1,4 @@
+from google.auth.transport.requests import Request
 from google.oauth2.credentials import Credentials
 from google_auth_oauthlib.flow import InstalledAppFlow
 from googleapiclient.discovery import build
@@ -20,8 +21,12 @@ class Authenticator:
 
     @classmethod
     def auth(cls):
-        msg = f'Doing auth for service {cls.SERVICE} ({cls.VERSION})...'
-        logger.info(msg)
+        """
+
+        Do initial authentication or refresh token if expired.
+
+        """
+        logger.info(f'Doing auth for service {cls.SERVICE} ({cls.VERSION})...')
 
         PATH_CREDS = cls.PATH / 'credentials.json'
         PATH_TOKEN = cls.PATH / f'{cls.SERVICE}.json'
@@ -29,9 +34,15 @@ class Authenticator:
         if PATH_TOKEN.exists():
             data_token = PATH_TOKEN.read_json()
             credentials = Credentials.from_authorized_user_info(data_token, cls.SCOPES)
+            if credentials.expired:
+                with logger.span(f'Credentials expired for {cls.SERVICE}. Will refresh...'):
+                    logger.warning(f'{cls.SERVICE}. {PATH_CREDS.exists()=} {PATH_TOKEN.exists()=} {credentials.valid=} {credentials.expired=} {credentials.expiry=}')
+                    credentials.refresh(Request())
+                    PATH_TOKEN.write_text(credentials.to_json())
         else:
             flow = InstalledAppFlow.from_client_secrets_file(PATH_CREDS, cls.SCOPES)
+            flow.authorization_url(prompt='consent', access_type='offline')
             credentials = flow.run_local_server(open_browser=False, port=cls.PORT)
             PATH_TOKEN.write_text(credentials.to_json())
-        service = build(cls.SERVICE, cls.VERSION, credentials=credentials)
-        return service
+
+        return build(cls.SERVICE, cls.VERSION, credentials=credentials)

fmtr/tools/http_tools.py

@@ -0,0 +1,26 @@
+import httpx
+from httpx_retries import RetryTransport
+
+from fmtr.tools import logging_tools
+
+logging_tools.logger.instrument_httpx()
+
+
+class Client(httpx.Client):
+    """
+
+    Instrumented client base
+
+    """
+
+    TRANSPORT = RetryTransport()
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, transport=self.TRANSPORT, **kwargs)
+
+
+client = Client()
+
+if __name__ == '__main__':
+    resp = client.get('http://httpbin.org/delay/10')
+    resp