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)
Ejemplo n.º 2
0
    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)