Commit dfef5726 authored by Tomas Krizek's avatar Tomas Krizek

dbhelper.DNSRepliesFactory: add serialization method

parent e8643dc4
......@@ -5,7 +5,8 @@ import os
import struct
import sys
import time
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Sequence # noqa
from typing import ( # noqa
Any, Callable, Dict, Iterator, List, Mapping, Optional, Tuple, Sequence)
import lmdb
......@@ -200,6 +201,7 @@ class DNSReply:
class DNSRepliesFactory:
"""Thread-safe factory to parse DNSReply objects from binary blob."""
def __init__(self, servers: Sequence[ResolverID]) -> None:
if not servers:
raise ValueError('One or more servers have to be specified')
......@@ -214,6 +216,19 @@ class DNSRepliesFactory:
logging.warning('Trailing data in buffer')
return replies
def serialize(self, replies: Mapping[ResolverID, DNSReply]) -> bytes:
data = []
for server in self.servers:
try:
reply = replies[server]
except KeyError:
raise ValueError('Missing reply for server "{}"!'.format(server))
else:
data.append(reply.binary)
if len(replies) > len(self.servers):
raise ValueError('Extra unexpected data to serialize!')
return b''.join(data)
class Database(ABC):
DB_NAME = b''
......
......@@ -88,7 +88,7 @@ def test_dns_reply_deserialization(binary, reply, remaining):
assert buff == remaining
def test_dns_replies_factory_init():
def test_dns_replies_factory():
with pytest.raises(ValueError):
DNSRepliesFactory([])
......@@ -97,13 +97,16 @@ def test_dns_replies_factory_init():
assert replies['a'] == DR_TIMEOUT
rf2 = DNSRepliesFactory(['a', 'b'])
replies = rf2.parse(DR_A_0_BIN + DR_ABCD_1_BIN)
bin_data = DR_A_0_BIN + DR_ABCD_1_BIN
replies = rf2.parse(bin_data)
assert replies['a'] == DR_A_0
assert replies['b'] == DR_ABCD_1
with pytest.raises(ValueError):
rf2.parse(DR_A_0_BIN + b'a')
assert rf2.serialize(replies) == bin_data
INT_3M = 3000000000
TIME_3M = 3000.0
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment