def test_ingest_input(self): with tempfile.NamedTemporaryFile('wb', prefix='dnsviz', delete=False) as example_bad_json: example_bad_json.write(b'{') with tempfile.NamedTemporaryFile('wb', prefix='dnsviz', delete=False) as example_no_version: example_no_version.write(b'{}') with tempfile.NamedTemporaryFile('wb', prefix='dnsviz', delete=False) as example_invalid_version_1: example_invalid_version_1.write(b'{ "_meta._dnsviz.": { "version": 1.11 } }') with tempfile.NamedTemporaryFile('wb', prefix='dnsviz', delete=False) as example_invalid_version_2: example_invalid_version_2.write(b'{ "_meta._dnsviz.": { "version": 5.0 } }') with gzip.open(EXAMPLE_AUTHORITATIVE, 'rb') as example_auth_in: with tempfile.NamedTemporaryFile('wb', prefix='dnsviz', delete=False) as example_auth_out: example_auth_out.write(example_auth_in.read()) try: args = ['-r', example_auth_out.name] arghelper = ArgHelper(self.logger) arghelper.build_parser('probe') arghelper.parse_args(args) arghelper.ingest_input() # Bad json args = ['-r', example_bad_json.name] arghelper = ArgHelper(self.logger) arghelper.build_parser('probe') arghelper.parse_args(args) with self.assertRaises(AnalysisInputError): arghelper.ingest_input() # No version args = ['-r', example_no_version.name] arghelper = ArgHelper(self.logger) arghelper.build_parser('probe') arghelper.parse_args(args) with self.assertRaises(AnalysisInputError): arghelper.ingest_input() # Invalid version args = ['-r', example_invalid_version_1.name] arghelper = ArgHelper(self.logger) arghelper.build_parser('probe') arghelper.parse_args(args) with self.assertRaises(AnalysisInputError): arghelper.ingest_input() # Invalid version args = ['-r', example_invalid_version_2.name] arghelper = ArgHelper(self.logger) arghelper.build_parser('probe') arghelper.parse_args(args) with self.assertRaises(AnalysisInputError): arghelper.ingest_input() finally: for tmpfile in (example_auth_out, example_bad_json, example_no_version, \ example_invalid_version_1, example_invalid_version_2): os.remove(tmpfile.name)
def test_ingest_names(self): args = ['example.com', 'example.net'] arghelper = ArgHelper(self.resolver, self.logger) arghelper.build_parser('probe', args) arghelper.ingest_names() self.assertEqual(list(arghelper.names), [ dns.name.from_text('example.com'), dns.name.from_text('example.net') ]) unicode_name = 'テスト' args = [unicode_name] arghelper = ArgHelper(self.resolver, self.logger) arghelper.build_parser('probe', args) arghelper.ingest_names() self.assertEqual(list(arghelper.names), [dns.name.from_text('xn--zckzah.')]) with tempfile.NamedTemporaryFile('wb', prefix='dnsviz', delete=False) as names_file: names_file.write('example.com\nexample.net\n'.encode('utf-8')) with tempfile.NamedTemporaryFile('wb', prefix='dnsviz', delete=False) as names_file_unicode: try: names_file_unicode.write( ('%s\n' % (unicode_name)).encode('utf-8')) # python3/python2 dual compatibility except UnicodeDecodeError: names_file_unicode.write(('%s\n' % (unicode_name))) with tempfile.NamedTemporaryFile('wb', prefix='dnsviz', delete=False) as example_names_only: example_names_only.write( b'{ "_meta._dnsviz.": { "version": 1.2, "names": [ "example.com.", "example.net.", "example.org." ] } }' ) try: args = ['-f', names_file.name] arghelper = ArgHelper(self.resolver, self.logger) arghelper.build_parser('probe', args) arghelper.ingest_names() self.assertEqual(list(arghelper.names), [ dns.name.from_text('example.com'), dns.name.from_text('example.net') ]) args = ['-f', names_file_unicode.name] arghelper = ArgHelper(self.resolver, self.logger) arghelper.build_parser('probe', args) arghelper.ingest_names() self.assertEqual(list(arghelper.names), [dns.name.from_text('xn--zckzah.')]) args = ['-r', example_names_only.name] arghelper = ArgHelper(self.resolver, self.logger) arghelper.build_parser('probe', args) arghelper.ingest_input() arghelper.ingest_names() self.assertEqual(list(arghelper.names), [ dns.name.from_text('example.com'), dns.name.from_text('example.net'), dns.name.from_text('example.org') ]) args = ['-r', example_names_only.name, 'example.com'] arghelper = ArgHelper(self.resolver, self.logger) arghelper.build_parser('probe', args) arghelper.ingest_input() arghelper.ingest_names() self.assertEqual(list(arghelper.names), [dns.name.from_text('example.com')]) finally: for tmpfile in (names_file, names_file_unicode, example_names_only): os.remove(tmpfile.name)