def bulk_search_cmd(args):
APIClient.API_BASE_URL = args.api_url
if args.input == '-':
input_file = read_stdin_lines()
elif os.path.exists(args.input):
input_file = open(args.input)
else:
show_error('%s not found\n' % args.input)
sys.exit(-1)
data_type = args.data_type
domain = args.domain
dns_type = args.type
ip = args.ip
host = args.host
if ip and not validate_ip(ip):
show_error('"%s" is not a valid IP address\n' % ip)
sys.exit(-1)
username = args.username
password = args.password
if not username:
username = read_line("Username:"******"Password:")
proxies = None
if args.proxy:
proxies = {'http': args.proxy, 'https': args.proxy}
client = DnsDBClient(proxies=proxies)
login(client, username, password)
output = get_output_file(args.output)
try:
for line in input_file:
line = line.strip()
if data_type == 'domain':
domain = line
elif data_type == 'ip':
ip = line
elif data_type == 'type':
dns_type = line
elif data_type == 'host':
host = line
start_time = datetime.datetime.now()
result = client.retrieve_dns(domain=domain,
host=host,
dns_type=dns_type,
ip=ip)
process_output(result, output,
OutputFormatter(args.json, args.csv, args.format),
args.max)
if args.verbose:
print('Running time: %s' %
(datetime.datetime.now() - start_time))
except Exception as e:
if args.debug:
traceback.print_exc()
show_error(str(e) + '\n')
finally:
output.close()
input_file.close()
def search_cmd(args):
APIClient.API_BASE_URL = args.api_url
username = args.username
password = args.password
domain = args.domain
host = args.host
dns_type = args.type
ip = args.ip
start = args.start
get_all = args.all
check_search_params(domain, host, ip)
if not username:
username = read_line("Username:"******"Password:")
proxies = None
if args.proxy:
proxies = {'http': args.proxy, 'https': args.proxy}
client = DnsDBClient(proxies=proxies)
login(client, username, password)
output = get_output_file(args.output)
start_time = datetime.datetime.now()
try:
if get_all:
result = client.retrieve_dns(domain=domain,
host=host,
dns_type=dns_type,
ip=ip)
else:
result = client.search_dns(domain=domain,
host=host,
dns_type=dns_type,
ip=ip,
start=start)
process_output(result, output,
OutputFormatter(args.json, args.csv, args.format),
args.max)
except Exception as e:
if isinstance(e, AuthenticationError):
os.remove(CACHE_PATH)
if args.debug:
traceback.print_exc()
show_error(str(e) + '\n')
finally:
output.flush()
if args.verbose:
print('Running time: %s' % (datetime.datetime.now() - start_time))
output.close()
def test_retrieve_dns_without_enough_remaining_request():
total = 200
mockserver.users = [User(default_username, default_password, 1)]
mockserver.max_page_size = 100
mockserver.dns_records = [
{
'host': 'a.com',
'type': 'a',
'value': '1.1.1.1'
},
] * total
client = DnsDBClient()
client.login(default_username, default_password)
results = client.retrieve_dns(domain='a.com')
assert_equal(len(results), total)
for _ in results:
pass
def test_search_dns():
user = User(default_username, default_password, 2)
mockserver.users = [user]
mockserver.dns_records = [
{
'host': 'a.com',
'type': 'a',
'value': '1.1.1.1'
},
] * 100
client = DnsDBClient()
client.login(default_username, default_password)
results = client.search_dns(domain='a.com',
ip='1.1.0.0',
dns_type='a',
host='c.a.com')
assert_equal(len(results), 30)
assert_equal(user.remaining_request, 1)
results = client.search_dns(domain='a.com', start=99)
assert_equal(len(results), 1)
assert_equal(user.remaining_request, 0)
for record in results:
assert_equal(record['host'], record.host)
assert_equal(record['type'], record.type)
assert_equal(record['value'], record.value)
data = json.loads(str(record))
assert_equal(data['host'], record.host)
assert_equal(data['type'], record.type)
assert_equal(data['value'], record.value)
def resources_cmd(args):
APIClient.API_BASE_URL = args.api_url
username = args.username
password = args.password
if not username:
username = read_line("Username:"******"Password:"******"Remaining DNS request: %s" % resources.remaining_dns_request)
if args.verbose:
print('Running time: %s' % (datetime.datetime.now() - start_time))
except Exception as e:
if args.debug:
traceback.print_exc()
show_error(str(e) + '\n')
def test_retrieve_dns():
total = 200
mockserver.users = [User(default_username, default_password, 2)]
mockserver.max_page_size = 100
mockserver.dns_records = [
{
'host': 'a.com',
'type': 'a',
'value': '1.1.1.1'
},
] * total
client = DnsDBClient()
client.login(default_username, default_password)
client.access_token.expire_at = datetime.datetime.now()
assert_true(client.access_token.has_expired())
results = client.retrieve_dns(domain='a.com',
dns_type='a',
ip='1.1.1.1',
host='c.a.com')
assert_equal(len(results), total)
count = 0
for _ in results:
count += 1
assert_equal(count, total)
def test_login():
mockserver.default_access_token = 'dnsdb api access-token'
client = DnsDBClient()
assert_false(client.is_login())
client.login(default_username, default_password)
assert_true(client.is_login())
assert_equal(
json.loads(str(client.access_token))['access_token'],
mockserver.default_access_token)
assert_equal(client.access_token.token, mockserver.default_access_token)
def test_search_dns_without_remaining_request():
mockserver.users = [User(default_username, default_password, 0)]
client = DnsDBClient()
client.login(default_username, default_password)
client.search_dns(domain='a.com')
def test_search_dns_without_params():
client = DnsDBClient()
client.login(default_username, default_password)
client.search_dns()
def test_search_dns_without_login():
client = DnsDBClient()
client.search_dns(domain='a.com')
def test_login_failed():
client = DnsDBClient()
client.login('test', '12345')
def test_get_resources_without_login():
client = DnsDBClient()
client.get_resources()
def test_get_resources():
mockserver.users = [User(default_username, default_password, 100)]
client = DnsDBClient()
client.login(default_username, default_password)
resources = client.get_resources()
assert_equal(resources.remaining_dns_request, 100)
def test_retrieve_dns_without_login():
client = DnsDBClient()
try:
client.retrieve_dns(domain='a.com')
except DnsDBException as e:
assert_equal(str(e), 'Require login')
def test_search_dns_return_504():
client = DnsDBClient()
client.login(default_username, default_password)
mockserver.return_error_response = GatewayTimeoutError
client.search_dns(domain='a.com')
def test_search_dns_return_500():
client = DnsDBClient()
client.login(default_username, default_password)
mockserver.return_error_response = InternalServerError
client.search_dns(domain='a.com')