-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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
- Loading branch information
Showing
9 changed files
with
231 additions
and
48 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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', | ||
] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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' |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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', | ||
] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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' |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters