#!/usr/bin/env python3
# -*- coding: utf-8 -*-

#    Unit test for DNSRecon's dnshelper library
#    Author: Filippo Lauria (@filippolauria)
#
#    Copyright (C) 2022  Carlos Perez
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; Applies version 2 of the License.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#    See the GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, write to the Free Software
#    Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

from dnsrecon.lib.dnshelper import DnsHelper
from netaddr import IPAddress
from re import match


class Test_Lib_dnshelper():

    def test_get_a(self):
        helper = DnsHelper('google.com')
        records = helper.get_a('ipv4.google.com')
        for record in records:
            assert record[0] in ['A', 'CNAME']

    def test_get_aaaa(self):
        helper = DnsHelper('google.com')
        records = helper.get_aaaa('ipv6.google.com')
        for record in records:
            assert record[0] in ['AAAA', 'CNAME']

    def test_get_mx(self):
        helper = DnsHelper('google.com')
        records = helper.get_mx()
        for record in records:
            assert record[0] == 'MX'

    def test_get_ip(self):
        helper = DnsHelper('google.com')
        records = helper.get_ip('google.com')
        for record in records:
            ip = IPAddress(record[2])
            assert ip.version in [4, 6]  # ~ redundant

    def test_get_txt(self):
        helper = DnsHelper('gmail.com')
        records = helper.get_txt()
        for record in records:
            assert record[0] == 'TXT'

    def test_get_ns(self):
        helper = DnsHelper('megacorpone.com')
        records = helper.get_ns()
        for record in records:
            assert record[0] == 'NS'

    def test_get_soa(self):
        helper = DnsHelper('megacorpone.com')
        records = helper.get_soa()
        for record in records:
            assert record[0] == 'SOA'

    def test_get_srv(self):
        helper = DnsHelper('nsztm1.digi.ninja')
        records = helper.get_srv('_sip._tcp.zonetransfer.me')
        for record in records:
            assert record[0] == 'SRV'

    def test_zone_transfer(self):
        helper = DnsHelper('megacorpone.com')
        records = helper.zone_transfer()
        assert len(records) == 30

    def test_get_ptr(self):
        helper = DnsHelper('megacorpone.com')
        records = helper.get_ptr('51.79.37.18')
        assert len(records) == 1 and match(r'^.+\.megacorpone\.com$', records[0][1])
