dataapi.py 1.61 KB
Newer Older
1 2 3 4
'''
Shared load/save methods for intermediate results
'''

5
import functools
6 7 8 9 10
import logging
import pickle
from typing import Dict, Set

import dns.name
11
from ednsevalzone import AnIPAddress
12

13
def load_nsname2ipset() -> Dict[dns.name.Name, Set[AnIPAddress]]:
14 15 16 17
    """raises FileNotFoundError"""
    logging.info('loading NS name -> IP address mapping')
    with open('nsname2ipset.pickle', 'rb') as nsname2ipset_pickle:
        nsname2ipset = pickle.load(nsname2ipset_pickle)
18
    ip_cnt = functools.reduce(lambda cnt, ipset: cnt + len(ipset), nsname2ipset.values(), 0)
19
    logging.info('loaded %s unique IP addresses for %s NS names',
20
                 ip_cnt, len(nsname2ipset))
21 22
    return nsname2ipset

23
def load_domain2ipset() -> Dict[dns.name.Name, Set[AnIPAddress]]:
24 25 26 27
    """raises FileNotFoundError"""
    logging.info('loading domain-IP sets')
    with open('domain2ipset.pickle', 'rb') as domain2ipset_pickle:
        domain2ipset = pickle.load(domain2ipset_pickle)
28 29
    return domain2ipset

30
def save_nsname2ipset(nsname2ipset: Dict[dns.name.Name, Set[AnIPAddress]]) -> None:
31
    if nsname2ipset:
32
        ip_cnt = functools.reduce(lambda cnt, ipset: cnt + len(ipset), nsname2ipset.values(), 0)
33 34 35 36 37 38 39 40 41 42 43 44 45 46 47
    else:
        ip_cnt = 0
    logging.info('pickling NS name -> IP address mapping '
                 '(%s unique IP addresses for %s NS names)',
                 ip_cnt, len(nsname2ipset))
    with open('nsname2ipset.pickle', 'wb') as nsname2ipset_pickle:
        pickle.dump(
            nsname2ipset,
            nsname2ipset_pickle
            )

class DataStore():
    def __init__(self):
        self.zone = None
        self.nsname2ipset = None