Adds support for adapters

Adapters can be used to customize the result loading and storing process by applying transforms to the results. This is used as an example in the certcheck plugin which deserializes the certificates from the result.

- Adds Adapter support
- Removes transform from utils

plugins/src/dnsmule_plugins/certcheck/__init__.py

@@ -1,6 +1,8 @@
 from typing import Optional, Callable, Collection
 
 from dnsmule import DNSMule, Domain, Plugin
+from dnsmule.adapter import patch_storage, Adapter
+from .adapter import load_result, save_result
 from .rule import CertChecker
 
 
@@ -13,6 +15,7 @@ def get_callback(self, mule: DNSMule) -> Optional[Callable[[Collection[Domain]],
             return mule.scan
 
     def register(self, mule: DNSMule):
+        patch_storage(mule.storage, Adapter(loader=load_result, saver=save_result))
         mule.rules.register(CertChecker.id)(CertChecker.creator(self.get_callback(mule)))
 
 

plugins/src/dnsmule_plugins/certcheck/adapter.py

@@ -0,0 +1,25 @@
+from . import certificates
+
+
+def load_result(result):
+    if 'resolvedCertificates' in result.data:
+        result.data['resolvedCertificates'] = {
+            certificates.Certificate.from_json(o)
+            for o in result.data['resolvedCertificates']
+        }
+    return result
+
+
+def save_result(result):
+    if 'resolvedCertificates' in result.data:
+        result.data['resolvedCertificates'] = [
+            certificates.Certificate.to_json(o)
+            for o in result.data['resolvedCertificates']
+        ]
+    return result
+
+
+__all__ = [
+    'load_result',
+    'save_result',
+]

plugins/src/dnsmule_plugins/certcheck/rule.py

@@ -1,7 +1,6 @@
 from typing import Callable, Collection, List, Optional
 
 from dnsmule import Rule, Result, Record
-from dnsmule.utils import extend_set, transform_set
 from . import certificates
 from .domains import process_domains
 
@@ -37,13 +36,7 @@ def __call__(self, record: Record) -> Result:
             )
         }
         if certs:
-            with transform_set(
-                    record.result.data,
-                    'resolvedCertificates',
-                    certificates.Certificate.from_json,
-                    certificates.Certificate.to_json
-            ):
-                extend_set(record.result.data, 'resolvedCertificates', certs)
+            record.result.data.setdefault('resolvedCertificates', set()).update(certs)
             if self.callback:
                 domains = [*process_domains(
                     domain

plugins/test/test_certcheck.py

@@ -77,8 +77,8 @@ def test_call_add_certs(mock_collection):
     assert len((check(r)).data['resolvedCertificates']) == 2, 'Failed to remove duplicates'
 
     certs = result.data['resolvedCertificates']
-    assert cert1.to_json() in certs, 'Failed to append existing data'
-    assert cert2.to_json() in certs, 'Failed to append data'
+    assert cert1 in certs, 'Failed to append existing data'
+    assert cert2 in certs, 'Failed to append data'
 
     assert len((check(r)).data['resolvedCertificates']) == 2, 'Failed to remove duplicates'
 

plugins/test/test_certcheck_adapter.py

@@ -0,0 +1,40 @@
+import datetime
+
+from dnsmule import Result, Domain, RRType
+from dnsmule_plugins.certcheck.adapter import save_result, load_result
+from dnsmule_plugins.certcheck.certificates import Certificate
+
+test_cert = Certificate(
+    version='2',
+    common='example.com',
+    issuer='test',
+    valid_from=datetime.datetime.now(),
+    valid_until=datetime.datetime.now(),
+    alts=[],
+)
+
+
+def test_saving():
+    result = Result(Domain('example.com'), RRType.TXT)
+    result.data['resolvedCertificates'] = {test_cert}
+    result = save_result(result)
+    assert result.data['resolvedCertificates'] == [test_cert.to_json()]
+
+
+def test_loading():
+    result = Result(Domain('example.com'), RRType.TXT)
+    result.data['resolvedCertificates'] = [test_cert.to_json()]
+    result = load_result(result)
+    assert result.data['resolvedCertificates'] == {test_cert}
+
+
+def test_saving_no_key():
+    result = Result(Domain('example.com'), RRType.TXT)
+    result = save_result(result)
+    assert 'resolvedCertificates' not in result.data, 'Added key when not present'
+
+
+def test_loading_no_key():
+    result = Result(Domain('example.com'), RRType.TXT)
+    result = load_result(result)
+    assert 'resolvedCertificates' not in result.data, 'Added key when not present'

src/dnsmule/adapter.py

@@ -0,0 +1,81 @@
+from abc import abstractmethod
+from types import MethodType
+from typing import Iterable, Optional
+
+from . import Result, Domain
+from .baseclasses import Identifiable
+from .storages import Query
+from .storages.abstract import Storage
+
+
+class Adapter(Identifiable):
+    """Adds capability to patch results coming and going from storage
+    """
+
+    def __init__(self, loader, saver):
+        """
+        :param loader: Function to call when loading a result, returns a result with modifications applied
+        :param saver: Function to call when saving a result, returns a result with modifications applied
+        """
+        self.loader = loader
+        self.saver = saver
+
+    @abstractmethod
+    def from_storage(self, result: Result) -> Result:
+        """Convert entries in data from a storage form to something user defined
+        """
+        return self.loader(result)
+
+    @abstractmethod
+    def to_storage(self, result: Result) -> Result:
+        """Convert to form suitable for going into the storage
+        """
+        return self.saver(result)
+
+
+def patch_method(instance: Storage, method: str):
+    """
+    Decorator to patch an instance method
+
+    **Note:** The usual *self* is actually the old method instance
+    """
+    old_method = getattr(instance, method)
+
+    def patcher(target):
+        setattr(instance, method, MethodType(target, old_method))
+        return target
+
+    return patcher
+
+
+def patch_storage(delegate: Storage, adapter: Adapter) -> None:
+    """
+    Patches a storage with the given adapter
+
+    This can be used to add transformations to data in result so that it is easier to handle in code.
+    A usual use-case would be to deserialize some json values from the result data.
+    """
+
+    @patch_method(delegate, 'store')
+    def _(method, result: Result) -> None:
+        return method(adapter.to_storage(result))
+
+    @patch_method(delegate, 'fetch')
+    def _(method, domain: Domain) -> Optional[Result]:
+        result = method(domain)
+        if result is not None:
+            return adapter.from_storage(result)
+
+    @patch_method(delegate, 'results')
+    def _(method) -> Iterable[Result]:
+        yield from map(adapter.from_storage, method())
+
+    @patch_method(delegate, 'query')
+    def _(method, query: Query) -> Iterable[Result]:
+        yield from map(adapter.from_storage, method(query))
+
+
+__all__ = [
+    'Adapter',
+    'patch_storage',
+]

src/dnsmule/utils/misc.py

@@ -1,6 +1,5 @@
-from contextlib import contextmanager
 from pathlib import Path
-from typing import Union, Callable, Iterator, TypeVar, Any, Dict, Tuple, Iterable
+from typing import Union, Iterator, TypeVar, Any, Dict, Tuple, Iterable
 
 K = TypeVar('K')
 V = TypeVar('V')
@@ -127,25 +126,10 @@ def extend_set(data: Dict[str, Any], key: str, values: Iterable[Any]):
     data[key] = target
 
 
-@contextmanager
-def transform_set(data: Dict[str, Any], key: str, function: Callable[[T], R], inverse: Callable[[R], T]) -> None:
-    """
-    Transforms a list based set using the given function
-
-    **Note:** Creates a new container for each transform
-
-    **Note:** Modifies data for the duration of the contextmanager
-    """
-    data[key] = [function(o) for o in data[key]] if key in data else []
-    yield data[key]
-    data[key] = [inverse(o) for o in data[key]] if key in data else []
-
-
 __all__ = [
     'load_data',
     'left_merge',
     'extend_set',
-    'transform_set',
     'empty',
     'join_values',
 ]

test/test_adapter.py

@@ -0,0 +1,77 @@
+from dnsmule import DNSMule
+from dnsmule.adapter import Adapter, patch_storage
+from dnsmule.storages import Query
+
+
+def add_loaded(result):
+    result.data['loaded'] = True
+    return result
+
+
+def add_stored(result):
+    result.data['stored'] = True
+    return result
+
+
+def test_patching_fetch(generate_result):
+    result = generate_result()
+    mule = DNSMule.make()
+    mule.storage.store(result)
+
+    patch_storage(mule.storage, Adapter(add_loaded, add_stored))
+    result = mule.storage.fetch(domain=result.domain)
+
+    assert result.data['loaded'], 'Failed to go through the adapter'
+
+
+def test_patching_fetch_no_result(generate_result):
+    mule = DNSMule.make()
+
+    patch_storage(mule.storage, Adapter(add_loaded, add_stored))
+    result = mule.storage.fetch(domain='example')
+
+    assert result is None, 'Produced something other than what was intended'
+
+
+def test_patching_results(generate_result):
+    result = generate_result()
+    mule = DNSMule.make()
+    mule.storage.store(result)
+
+    patch_storage(mule.storage, Adapter(add_loaded, add_stored))
+    result = next(mule.storage.results())
+
+    assert result.data['loaded'], 'Failed to go through the adapter'
+
+
+def test_patching_query(generate_result):
+    result = generate_result()
+    mule = DNSMule.make()
+    mule.storage.store(result)
+
+    patch_storage(mule.storage, Adapter(add_loaded, add_stored))
+    result = next(mule.storage.query(query=Query(domains=[result.domain])))
+
+    assert result.data['loaded'], 'Failed to go through the adapter'
+
+
+def test_patching_store(generate_result):
+    result = generate_result()
+    mule = DNSMule.make()
+    old_method = mule.storage.fetch
+
+    patch_storage(mule.storage, Adapter(add_loaded, add_stored))
+    mule.storage.store(result)
+
+    result = old_method(result.domain)
+    assert result.data['stored'], 'Failed to go through the adapter on store'
+
+
+def test_adapter_calls_correctly(generate_result):
+    save = object()
+    load = object()
+
+    adapter = Adapter(lambda _: load, lambda _: save)
+
+    assert adapter.from_storage(generate_result()) is load, 'Called the wrong method on load'
+    assert adapter.to_storage(generate_result()) is save, 'Called the wrong method on save'

test/utils/test_utils.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-from dnsmule.utils import load_data, left_merge, extend_set, transform_set, join_values, empty
+from dnsmule.utils import load_data, left_merge, extend_set, join_values, empty
 
 
 def test_empty_raises():
@@ -116,23 +116,3 @@ def test_extend_set_de_duplicates_new():
     store = {'key': ['a', 'a', 'a']}
     extend_set(store, 'key', ['a', 'a', 'a'])
     assert store['key'] == ['a'], 'Failed to de-duplicate'
-
-
-def test_transform_adds_value():
-    store = {}
-    with transform_set(store, 'key', int, str):
-        assert 'key' in store, 'Failed to add value'
-
-
-def test_transform_creates_new_list():
-    target = ['1', '2', '3']
-    store = {'key': target}
-    with transform_set(store, 'key', int, str):
-        assert store['key'] is not target
-
-
-def test_transform_values():
-    store = {'key': ['1', '2', '3']}
-    with transform_set(store, 'key', int, str):
-        assert store['key'] == [1, 2, 3]
-    assert store['key'] == ['1', '2', '3']