1

import itertools

     2

import re

     3

import unittest

     4

     5

from pvl.dns import zone, process

     6

from StringIO import StringIO

     7

     8

class File(StringIO):

     9

    @classmethod

    10

    def lines (cls, *lines):

    11

        return cls('\n'.join(lines) + '\n')

    12

    13

    def __init__(self, buffer, name='test.file'):

    14

        StringIO.__init__(self, buffer)

    15

        self.name = name

    16

    17

class TestMixin(object):

    18

    def assertEqualWhitespace(self, value, expected):

    19

        # normalize

    20

        value = re.sub(r'\s+', ' ', value)

    21

    22

        self.assertEqual(value, expected)

    23

    24

    def assertZoneLineEqual(self, zl, expected):

    25

        self.assertEqual(unicode(zl), expected)

    26

    27

    def assertZoneLinesEqual(self, zls, expected):

    28

        for zl, expect in itertools.izip_longest(zls, expected):

    29

            self.assertZoneLineEqual(zl, expect)

    30

    31

    def assertZoneItemEqual(self, zr, expect):

    32

        self.assertEqualWhitespace(unicode(zr), expect)

    33

    34

    def assertZoneEqual(self, zrs, expected):

    35

        for zr, expect in itertools.izip_longest(zrs, expected):

    36

            self.assertZoneItemEqual(zr, expect)

    37

    38

    assertZoneRecordEqual = assertZoneItemEqual

    39

    assertZoneRecordsEqual = assertZoneEqual

    40

    41

class ZoneLineTest(TestMixin, unittest.TestCase):

    42

    def testZoneLine(self):

    43

        self.assertEqual(unicode(zone.ZoneLine(['foo', 'A', '192.0.2.1'])), "foo\tA\t192.0.2.1")

    44

    45

    def testZoneLineParse(self):

    46

        self.assertZoneLinesEqual(zone.ZoneLine.parse(File("foo A 192.0.2.1")), ["foo\tA\t192.0.2.1"])

    47

    48

    def testZoneLineParseWhitespace(self):

    49

        self.assertZoneLinesEqual(zone.ZoneLine.parse(File("foo \t A\t192.0.2.1")), ["foo\tA\t192.0.2.1"])

    50

    51

    def testZoneLineParseComment(self):

    52

        self.assertZoneLinesEqual(zone.ZoneLine.parse(File("foo A 192.0.2.1   ;   bar")), ["foo\tA\t192.0.2.1\t; bar"])

    53

    54

    def testZoneLineParseQuoteSimple(self):

    55

        self.assertZoneLinesEqual(zone.ZoneLine.parse(File("foo TXT \"asdf quux\"")), ["foo\tTXT\tasdf quux"])

    56

    57

    def testZoneLineParseQuoteTrailing(self):

    58

        self.assertZoneLinesEqual(zone.ZoneLine.parse(File("foo TXT \"asdf quux\" ok ; yay")), ["foo\tTXT\tasdf quux\tok\t; yay"])

    59

    60

    def testZoneLineParseIndent(self):

    61

        self.assertZoneLinesEqual(zone.ZoneLine.parse(File(" A 192.0.2.2")), ["\tA\t192.0.2.2"])

    62

    63

    def testZoneLineParseMultilineSingle(self):

    64

        self.assertZoneLinesEqual(zone.ZoneLine.parse(File("@ SOA ( 1 2 3 4 5 )")), ["@\tSOA\t1\t2\t3\t4\t5"])

    65

    66

    def testZoneLineParseMultiline(self):

    67

        self.assertZoneLinesEqual(zone.ZoneLine.parse(File.lines("@ SOA (", "\t1\t", "2", "\t3\t\t4", "\t5 )")), ["@\tSOA\t1\t2\t3\t4\t5"])

    68

    69

    def testZoneLineLoad(self):

    70

        self.assertZoneRecordsEqual(

    71

            zone.ZoneLine.load(File.lines(

    72

                "$TTL 3600",

    73

                " ",

    74

                "@ NS ns1",

    75

                "  NS ns2",

    76

                "$ORIGIN asdf", # note lack of .

    77

                "foo A 192.0.2.1",

    78

                "$ORIGIN quux.test.", # note lack of .

    79

                "bar A 192.0.2.2",

    80

            )), [

    81

                "$TTL 3600",

    82

                "@ NS ns1",

    83

                " NS ns2",

    84

                "$ORIGIN asdf",

    85

                "foo A 192.0.2.1",

    86

                "$ORIGIN quux.test.",

    87

                "bar A 192.0.2.2",

    88

            ]

    89

        )

    90

    91

class ZoneDirectiveTest(TestMixin, unittest.TestCase):

    92

    def testZoneDirective(self):

    93

        self.assertEqual(unicode(zone.ZoneDirective('ORIGIN', ['foo.'])), "$ORIGIN\tfoo.")

    94

    95

    def testZoneDirectiveBuild(self):

    96

        self.assertEqual(unicode(zone.ZoneDirective.build('ORIGIN', 'foo.')), "$ORIGIN\tfoo.")

    97

    98

    def testZoneDirectiveParse(self):

    99

        self.assertEqual(unicode(zone.ZoneDirective.parse(['$ORIGIN', 'foo.'])), "$ORIGIN\tfoo.")

   100

   101

    def testZoneDirectiveParseUpper(self):

   102

        self.assertEqual(unicode(zone.ZoneDirective.parse(['$include', 'foo.zone'])), "$INCLUDE\tfoo.zone")

   103

   104

    def testZoneDirectiveComment(self):

   105

        self.assertEqual(unicode(zone.ZoneDirective('ORIGIN', ['foo.'], comment="bar")), "$ORIGIN\tfoo.\t; bar")

   106

   107

class ZoneRecordTest(TestMixin, unittest.TestCase):

   108

    def testZoneRecordShort(self):

   109

        rr = zone.ZoneRecord('test', None, None, 'A', ['192.0.2.1'])

   110

   111

        self.assertZoneRecordEqual((rr), 'test A 192.0.2.1')

   112

   113

    def testZoneRecordImplicit(self):

   114

        rr = zone.ZoneRecord(None, None, None, 'A', ['192.0.2.1'])

   115

   116

        self.assertZoneRecordEqual((rr), ' A 192.0.2.1')

   117

   118

    def testZoneRecordFull(self):

   119

        rr = zone.ZoneRecord('test', 60, 'IN', 'A', ['192.0.2.1'])

   120

   121

        self.assertZoneRecordEqual((rr), 'test 60 IN A 192.0.2.1')

   122

   123

    def testZoneRecordComment(self):

   124

        rr = zone.ZoneRecord('test', None, None, 'A', ['192.0.2.1'], comment='Testing')

   125

   126

        self.assertZoneRecordEqual((rr), 'test A 192.0.2.1 ; Testing')

   127

   128

    def testZoneRecordA(self):

   129

        self.assertZoneRecordEqual((zone.ZoneRecord.A('test', '192.0.2.1')), "test A 192.0.2.1")

   130

   131

    def testZoneRecordAAAA(self):

   132

        self.assertZoneRecordEqual((zone.ZoneRecord.AAAA('test', '2001:db8::c000:201')), "test AAAA 2001:db8::c000:201")

   133

   134

    def testZoneRecordCNAME(self):

   135

        self.assertZoneRecordEqual((zone.ZoneRecord.CNAME('test', 'test.example.net.')), "test CNAME test.example.net.")

   136

   137

    def testZoneRecordTXT(self):

   138

        self.assertZoneRecordEqual((zone.ZoneRecord.TXT('test', u"Foo Bar")), u"test TXT \"Foo Bar\"")

   139

   140

    def testZoneRecordPTR(self):

   141

        self.assertZoneRecordEqual(zone.ZoneRecord.PTR('1', 'test.example.com.'), "1 PTR test.example.com.")

   142

   143

    def testZoneRecordMX(self):

   144

        self.assertZoneRecordEqual(zone.ZoneRecord.MX('@', 10, 'mail1'), "@ MX 10 mail1")

   145

   146

    def testZoneRecordBuildTTL(self):

   147

        self.assertZoneRecordEqual(zone.ZoneRecord.A('test', '192.0.2.1', ttl=60), "test 60 A 192.0.2.1")

   148

   149

    def testZoneRecordBuildZeroTTL(self):

   150

        self.assertZoneRecordEqual(zone.ZoneRecord.A('test', '192.0.2.1', ttl=0), "test 0 A 192.0.2.1")

   151

   152

    def testZoneRecordBuildImplicit(self):

   153

        self.assertZoneRecordEqual(zone.ZoneRecord.A(None, '192.0.2.1'), " A 192.0.2.1")

   154

   155

    def testZoneRecordBuildCls(self):

   156

        self.assertZoneRecordEqual(zone.ZoneRecord.A('foo', '192.0.2.1', cls='in'), "foo IN A 192.0.2.1")

   157

   158

    def testZoneRecordParse(self):

   159

        self.assertZoneRecordEqual(zone.ZoneRecord.parse(['test', 'A', '192.0.2.1']), "test A 192.0.2.1")

   160

        self.assertZoneRecordEqual(zone.ZoneRecord.parse(['', 'A', '192.0.2.1']), " A 192.0.2.1")

   161

        self.assertZoneRecordEqual(zone.ZoneRecord.parse(['test', '60', 'A', '192.0.2.1']), "test 60 A 192.0.2.1")

   162

        self.assertZoneRecordEqual(zone.ZoneRecord.parse(['test', '60', 'in', 'A', '192.0.2.1']), "test 60 IN A 192.0.2.1")

   163

   164

    def testZoneRecordParseError(self):

   165

        with self.assertRaises(zone.ZoneLineError):

   166

            self.assertZoneRecordEqual(zone.ZoneRecord.parse(['test', 'A']), None)

   167

   168

    def testZoneRecordLoad(self):

   169

        self.assertZoneRecordsEqual(

   170

            zone.ZoneRecord.load(File.lines(

   171

                "$TTL 3600",

   172

                " ",

   173

                "@ NS ns1",

   174

                "  NS ns2",

   175

                "$ORIGIN asdf", # relative

   176

                "foo A 192.0.2.1",

   177

                "$ORIGIN quux.test.", # absolute

   178

                "bar A 192.0.2.2",

   179

            ), 'test'), [

   180

                "@ 3600 NS ns1",

   181

                "@ 3600 NS ns2",

   182

                "foo 3600 A 192.0.2.1", # asdf.test

   183

                "bar 3600 A 192.0.2.2", # quux.test

   184

            ]

   185

        )

   186

   187

class TestProcessZone(TestMixin, unittest.TestCase):

   188

    def testZoneRecordLoad(self):

   189

        rrs = zone.ZoneLine.load(File("""

   190

$TTL 3600

   191

@                   SOA     foo.test. hostmaster.test. (

   192

                            0               ; serial

   193

                            1d              ; refresh

   194

                            5m              ; retry

   195

                            10d             ; expiry

   196

                            300             ; negative

   197

                    )

   198

   199

                    NS      foo

   200

                    NS      bar

   201

   202

foo                 A       192.0.2.1

   203

bar                 A       192.0.2.2

   204

   205

$INCLUDE "includes/test"

   206

"""))

   209

        rrs = list(process.zone_serial(rrs, 1337))

   211

   212

        self.assertZoneEqual(rrs, [

   213

            "$TTL 3600",

   214

            "@ SOA foo.test. hostmaster.test. 1337 1d 5m 10d 300",

   215

            " NS foo",

   216

            " NS bar",

   217

            "foo A 192.0.2.1",

   218

            "bar A 192.0.2.2",

   219

            "$INCLUDE .../includes/test",

   220

        ])

   224