File: rbldnsd.py

package info (click to toggle)
rbldnsd 0.998b~pre1-1
  • links: PTS, VCS
  • area: main
  • in suites: buster, stretch
  • size: 744 kB
  • ctags: 1,132
  • sloc: ansic: 8,212; python: 549; sh: 502; makefile: 249; awk: 33
file content (163 lines) | stat: -rw-r--r-- 5,067 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
""" Run an rbldnsd and send it DNS queries.


"""
from itertools import count
import subprocess
from tempfile import NamedTemporaryFile, TemporaryFile
import time
import unittest

try:
    import DNS
except ImportError:
    raise RuntimeError("The pydns library is not installed")

DUMMY_ZONE_HEADER = """
$SOA 0 example.org. hostmaster.example.com. 0 1h 1h 2d 1h
$NS 1d ns0.example.org
"""

class ZoneFile(object):
    def __init__(self, lines=None, no_header=False):
        self._file = NamedTemporaryFile()
        if not no_header:
            self._file.write(DUMMY_ZONE_HEADER)
        if lines is not None:
            self.writelines(lines)
        self._file.flush()

    @property
    def name(self):
        return self._file.name

    def write(self, str):
        self._file.write(str)
        self._file.flush()

    def writelines(self, lines):
        self._file.writelines("%s\n" % line for line in lines)
        self._file.flush()

class DaemonError(Exception):
    """ Various errors having to do with the execution of the daemon.
    """

class QueryRefused(Exception):
    """ Query to rbldnsd was REFUSED.
    """


class Rbldnsd(object):
    def __init__(self, datasets=None,
                 daemon_addr='localhost', daemon_port=5300,
                 daemon_bin='./rbldnsd',
                 stderr=None):
        self._daemon = None
        self.datasets = []
        self.daemon_addr = daemon_addr
        self.daemon_port = daemon_port
        self.daemon_bin = daemon_bin
        self.stderr = stderr

    def add_dataset(self, ds_type, file, soa='example.com'):
        self.datasets.append((soa, ds_type, file))

    def __enter__(self):
        self._start_daemon()
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        self._stop_daemon()

    def __del__(self):
        if self._daemon:
            self._stop_daemon()

    def query(self, name, qtype='TXT'):
        if not self._daemon:
            raise DaemonError("daemon not running")
        elif self._daemon.poll() is not None:
            raise DaemonError("daemon has died with code %d"
                              % self._daemon.returncode)

        req = DNS.Request(name=name, qtype=qtype, rd=0)
        resp = req.req(server=self.daemon_addr, port=self.daemon_port)
        status = resp.header['status']
        if status == 'REFUSED':
            raise QueryRefused("REFUSED")
        elif status == 'NXDOMAIN':
            return None
        else:
            assert status == 'NOERROR'
            assert len(resp.answers) == 1
            assert len(resp.answers[0]['data']) == 1
            return resp.answers[0]['data'][0]

    def _start_daemon(self):
        if len(self.datasets) == 0:
            raise ValueError("no datasets defined")

        cmd = [ self.daemon_bin, '-n',
                '-b', '%s/%u' % (self.daemon_addr, self.daemon_port),
                ]
        for zone, ds_type, file in self.datasets:
            if isinstance(file, basestring):
                filename = file
            else:
                filename = file.name
            cmd.append("%s:%s:%s" % (zone, ds_type, filename))

        self._stdout = TemporaryFile()
        self._daemon = daemon = subprocess.Popen(cmd, stdout=self._stdout,
                                                 stderr=self.stderr)

        # wait for rbldnsd to start responding
        time.sleep(0.1)
        for retry in count():
            if daemon.poll() is not None:
                raise DaemonError(
                    "rbldsnd exited unexpectedly with return code %d"
                    % daemon.returncode)
            try:
                self.query('dummy.nonexisting.zone')
                break
            except QueryRefused:
                break
            except DNS.DNSError as ex:
                if str(ex) != 'no working nameservers found':
                    raise
                elif retries > 10:
                    raise DaemonError(
                        "rbldnsd does not seem to be responding")

    def _stop_daemon(self):
        daemon = self._daemon

        if daemon.poll() is None:
            daemon.terminate()
            retries = count()
            while daemon.poll() is None:
                retry = next(retries)
                if retry == 30:
                    daemon.kill()
                elif retry == 50:
                    raise DaemonError("can not kill stop rbldnsd")
                time.sleep(0.1)

        self._daemon = None
        if daemon.returncode != 0:
            raise DaemonError("rbldnsd exited with code %d"
                              % daemon.returncode)

class TestRbldnsd(unittest.TestCase):
    def test(self):
        rbldnsd = Rbldnsd()
        test_zone = ZoneFile(lines=["1.2.3.4 :1: Success"])
        rbldnsd.add_dataset('ip4set', test_zone)
        with rbldnsd:
            self.assertEqual(rbldnsd.query('4.3.2.1.example.com'), 'Success')
            self.assertEqual(rbldnsd.query('5.3.2.1.example.com'), None)

if __name__ == '__main__':
    unittest.main()