[3] | 1 | """ |
---|
| 2 | Extension functions to help query/assert name service information. |
---|
| 3 | |
---|
| 4 | Functions: |
---|
| 5 | |
---|
| 6 | * dns_resolves -- assert that a host resolves to a specific IP address. |
---|
| 7 | * dns_a -- assert that a host directly resolves to a specific IP address |
---|
| 8 | * dns_cname -- assert that a host is an alias for another hostname. |
---|
| 9 | * dnx_mx -- assert that a given host is a mail exchanger for the given name. |
---|
| 10 | * dns_ns -- assert that a given hostname is a name server for the given name. |
---|
| 11 | """ |
---|
| 12 | |
---|
| 13 | import socket |
---|
| 14 | from twill.errors import TwillAssertionError |
---|
| 15 | |
---|
| 16 | try: |
---|
| 17 | import dns.resolver |
---|
| 18 | except ImportError: |
---|
| 19 | raise Exception("ERROR: must have dnspython installed to use the DNS extension module") |
---|
| 20 | |
---|
| 21 | def dns_a(host, ipaddress, server=None): |
---|
| 22 | """ |
---|
| 23 | >> dns_a <name> <ipaddress> [<name server>] |
---|
| 24 | |
---|
| 25 | Assert that <name> resolves to <ipaddress> (and is an A record). |
---|
| 26 | Optionally use the given name server. |
---|
| 27 | """ |
---|
| 28 | if not is_ip_addr(ipaddress): |
---|
| 29 | raise Exception("<ipaddress> parameter must be an IP address, not a hostname") |
---|
| 30 | |
---|
| 31 | for answer in _query(host, 'A', server): |
---|
| 32 | if ipaddress == answer.address: |
---|
| 33 | return True |
---|
| 34 | |
---|
| 35 | raise TwillAssertionError |
---|
| 36 | |
---|
| 37 | def dns_cname(host, cname, server=None): |
---|
| 38 | """ |
---|
| 39 | >> dns_cname <name> <alias_for> [<name server>] |
---|
| 40 | |
---|
| 41 | Assert that <name> is a CNAME alias for <alias_for> Optionally use |
---|
| 42 | <name server>. |
---|
| 43 | """ |
---|
| 44 | if is_ip_addr(cname): |
---|
| 45 | raise Exception("<alias_for> parameter must be a hostname, not an IP address") |
---|
| 46 | |
---|
| 47 | cname = dns.name.from_text(cname) |
---|
| 48 | |
---|
| 49 | for answer in _query(host, 'CNAME', server): |
---|
| 50 | if cname == answer.target: |
---|
| 51 | return True |
---|
| 52 | |
---|
| 53 | raise TwillAssertionError |
---|
| 54 | |
---|
| 55 | def dns_resolves(host, ipaddress, server=None): |
---|
| 56 | """ |
---|
| 57 | >> dns_resolves <name> <name2/ipaddress> [<name server>] |
---|
| 58 | |
---|
| 59 | Assert that <name> ultimately resolves to the given IP address (or |
---|
| 60 | the same IP address that 'name2' resolves to). Optionally use the |
---|
| 61 | given name server. |
---|
| 62 | """ |
---|
| 63 | if not is_ip_addr(ipaddress): |
---|
| 64 | ipaddress = _resolve_name(ipaddress, server) |
---|
| 65 | |
---|
| 66 | for answer in _query(host, 1, server): |
---|
| 67 | if ipaddress == answer.address: |
---|
| 68 | return True |
---|
| 69 | |
---|
| 70 | raise TwillAssertionError |
---|
| 71 | |
---|
| 72 | def dns_mx(host, mailserver, server=None): |
---|
| 73 | """ |
---|
| 74 | >> dns_mx <name> <mailserver> [<name server>] |
---|
| 75 | |
---|
| 76 | Assert that <mailserver> is a mailserver for <name>. |
---|
| 77 | """ |
---|
| 78 | mailserver = dns.name.from_text(mailserver) |
---|
| 79 | |
---|
| 80 | for rdata in _query(host, 'MX', server): |
---|
| 81 | if mailserver == rdata.exchange: |
---|
| 82 | return True |
---|
| 83 | |
---|
| 84 | raise TwillAssertionError |
---|
| 85 | |
---|
| 86 | def dns_ns(host, query_ns, server=None): |
---|
| 87 | """ |
---|
| 88 | >> dns_ns <domain> <nameserver> [<name server to use>] |
---|
| 89 | |
---|
| 90 | Assert that <nameserver> is a mailserver for <domain>. |
---|
| 91 | """ |
---|
| 92 | query_ns = dns.name.from_text(query_ns) |
---|
| 93 | |
---|
| 94 | for answer in _query(host, 'NS', server): |
---|
| 95 | if query_ns == answer.target: |
---|
| 96 | return True |
---|
| 97 | |
---|
| 98 | raise TwillAssertionError |
---|
| 99 | |
---|
| 100 | ### |
---|
| 101 | |
---|
| 102 | def is_ip_addr(text): |
---|
| 103 | """ |
---|
| 104 | Check the 'name' to see if it's just an IP address. |
---|
| 105 | """ |
---|
| 106 | |
---|
| 107 | try: |
---|
| 108 | v = dns.ipv4.inet_aton(text) |
---|
| 109 | return True |
---|
| 110 | except socket.error: |
---|
| 111 | return False |
---|
| 112 | |
---|
| 113 | def _resolve_name(name, server): |
---|
| 114 | """ |
---|
| 115 | Resolve the given name to an IP address. |
---|
| 116 | """ |
---|
| 117 | if is_ip_addr(name): |
---|
| 118 | return name |
---|
| 119 | |
---|
| 120 | r = dns.resolver.Resolver() |
---|
| 121 | if server: |
---|
| 122 | r.nameservers = [_resolve_name(server, None)] |
---|
| 123 | |
---|
| 124 | answers = r.query(name) |
---|
| 125 | |
---|
| 126 | answer = None |
---|
| 127 | for answer in answers: # @CTB !? |
---|
| 128 | break |
---|
| 129 | |
---|
| 130 | assert answer |
---|
| 131 | return str(answer) |
---|
| 132 | |
---|
| 133 | def _query(query, query_type, server): |
---|
| 134 | """ |
---|
| 135 | Query, perhaps via the given name server. (server=None to use default). |
---|
| 136 | """ |
---|
| 137 | r = dns.resolver.Resolver() |
---|
| 138 | if server: |
---|
| 139 | r.nameservers = [_resolve_name(server, None)] |
---|
| 140 | |
---|
| 141 | return r.query(query, query_type) |
---|