def test_info_set_keep_uniq_id(self):
        #
        # Create a new InfoSet, load it from the KB, confirm that it has
        # the same uniq_id
        #
        vuln = MockVuln(name='Foos')

        info_set_a, created = kb.append_uniq_group('a', 'b', vuln,
                                                   group_klass=MockInfoSetNames)

        self.assertTrue(created)

        info_set_b = kb.get('a', 'b')[0]

        self.assertEqual(info_set_a.get_uniq_id(),
                         info_set_b.get_uniq_id())

        #
        # Change the InfoSet a little bit by adding a new Info. That should
        # change the uniq_id
        #
        vuln = MockVuln(name='Foos')
        _, created = kb.append_uniq_group('a', 'b', vuln,
                                          group_klass=MockInfoSetNames)

        self.assertFalse(created)

        info_set_b = kb.get('a', 'b')[0]

        self.assertNotEqual(info_set_a.get_uniq_id(),
                            info_set_b.get_uniq_id())

 def test_pickleable_vuln(self):
     original_vuln = MockVuln()
     
     kb.append('a', 'b', original_vuln)
     unpickled_vuln = kb.get('a', 'b')[0]
     
     self.assertEqual(original_vuln, unpickled_vuln)

    def test_alert_if_target_is_301_all_internal_redir(self):
        """
        Tests that no info is created if the site redirects internally
        """
        core = w3afCore()

        httpretty.register_uri(httpretty.GET,
                               re.compile("w3af.com/(.*)"),
                               body='301',
                               status=301,
                               adding_headers={'Location': 'http://w3af.com/xyz'})

        target = core.target.get_options()
        target['target'].set_value('http://w3af.com/')
        core.target.set_options(target)

        core.plugins.set_plugins(['sqli'], 'audit')
        core.plugins.init_plugins()

        core.verify_environment()
        core.scan_start_hook()

        strategy = CoreStrategy(core)
        strategy.start()

        infos = kb.get('core', 'core')
        self.assertEqual(len(infos), 0, infos)

 def test_pickleable_info(self):
     original_info = MockInfo()
     
     kb.append('a', 'b', original_info)
     unpickled_info = kb.get('a', 'b')[0]
     
     self.assertEqual(original_info, unpickled_info)

 def test_return_all_for_plugin(self):
     i1 = MockInfo()
     i2 = MockInfo()
     i3 = MockInfo()
     
     kb.append('a', 'b', i1)
     kb.append('a', 'b', i2)
     kb.append('a', 'b', i3)
     
     self.assertEqual(kb.get('a', 'b'), [i1, i2, i3])

        def multi_append():
            for i in xrange(InfoSet.MAX_INFO_INSTANCES * 2):
                vuln = MockVuln()
                kb.append_uniq_group('a', 'b', vuln, group_klass=MockInfoSetTrue)

            info_set_list = kb.get('a', 'b')

            self.assertEqual(len(info_set_list), 1)

            info_set = info_set_list[0]
            self.assertEqual(len(info_set.infos), InfoSet.MAX_INFO_INSTANCES)
            return True

 def test_append(self):
     i1 = MockInfo()
     i2 = MockInfo()
     i3 = MockInfo()
     
     kb.append('a', 'b', i1)
     kb.append('a', 'b', i1)
     kb.append('a', 'b', i1)
     kb.append('a', 'b', i2)
     kb.append('a', 'b', i3)
     
     self.assertEqual(kb.get('a', 'b'), [i1, i1, i1, i2, i3])

    def test_append_uniq_var_bug_10Dec2012(self):
        i1 = MockInfo()
        i1.set_uri(URL('http://moth/abc.html'))
        i1.set_var('id')

        i2 = MockInfo()
        i2.set_uri(URL('http://moth/abc.html'))
        i2.set_var('id')

        kb.append_uniq('a', 'b', i1)
        kb.append_uniq('a', 'b', i2)
        self.assertEqual(kb.get('a', 'b'), [i1, ])

    def test_append_uniq_var_default(self):
        i1 = MockInfo()
        i1.set_uri(URL('http://moth/abc.html?id=1'))
        i1.set_dc(QueryString([('id', ['1'])]))
        i1.set_token(('id', 0))

        i2 = MockInfo()
        i2.set_uri(URL('http://moth/abc.html?id=3'))
        i2.set_dc(QueryString([('id', ['3'])]))
        i2.set_token(('id', 0))

        kb.append_uniq('a', 'b', i1)
        kb.append_uniq('a', 'b', i2)
        self.assertEqual(kb.get('a', 'b'), [i1, ])

 def test_save_append(self):
     """
     Although calling raw_write and then append is highly discouraged,
     someone would want to use it.
     """
     i0 = MockInfo()
     self.assertRaises(TypeError, kb.raw_write, 'a', 'b', i0)
     
     i1 = MockInfo()
     i2 = MockInfo()
     kb.append('a', 'b', i1)
     kb.append('a', 'b', i2)
     
     self.assertEqual(kb.get('a', 'b'), [i1, i2])

    def test_append_uniq_url_different(self):
        i1 = MockInfo()
        i1.set_uri(URL('http://moth/abc.html?id=1'))
        i1.set_dc(QueryString([('id', ['1'])]))
        i1.set_token(('id', 0))

        i2 = MockInfo()
        i2.set_uri(URL('http://moth/def.html?id=3'))
        i2.set_dc(QueryString([('id', ['3'])]))
        i2.set_token(('id', 0))

        kb.append_uniq('a', 'b', i1, filter_by='URL')
        kb.append_uniq('a', 'b', i2, filter_by='URL')
        self.assertEqual(kb.get('a', 'b'), [i1, i2])

    def test_append_uniq_url_uniq(self):
        i1 = MockInfo()
        i1.set_uri(URL('http://moth/abc.html?id=1'))
        i1.set_dc(QueryString([('id', '1')]))
        i1.set_var('id')

        i2 = MockInfo()
        i2.set_uri(URL('http://moth/abc.html?id=3'))
        i2.set_dc(QueryString([('id', '3')]))
        i2.set_var('id')

        kb.append_uniq('a', 'b', i1, filter_by='URL')
        kb.append_uniq('a', 'b', i2, filter_by='URL')
        self.assertEqual(kb.get('a', 'b'), [i1,])

 def test_store_in_kb(self):
     dt = DAVTemplate()
     dt.store_in_kb()
     
     stored_data = kb.get(*dt.get_kb_location())
     
     self.assertEqual(len(stored_data), 1)
     
     stored_vuln = stored_data[0]
     created_vuln = dt.create_vuln()
     
     stored_vuln.set_id(created_vuln.get_id())
     
     self.assertEqual(stored_vuln, created_vuln)
     
     

 def test_pickleable_shells(self):
     pool = Pool(1)
     xurllib = ExtendedUrllib()
     
     original_shell = Shell(MockVuln(), xurllib, pool)
     
     kb.append('a', 'b', original_shell)
     unpickled_shell = kb.get('a', 'b')[0]
     
     self.assertEqual(original_shell, unpickled_shell)
     self.assertEqual(unpickled_shell.worker_pool, None)
     self.assertEqual(unpickled_shell._uri_opener, None)
     
     pool.terminate()
     pool.join()
     xurllib.end()

    def test_append_uniq_var_not_uniq_diff_token_name(self):
        i1 = MockInfo()
        i1.set_uri(URL('http://moth/abc.html?id=1&foo=bar'))
        i1.set_dc(QueryString([('id', ['1']),
                               ('foo', ['bar'])]))
        i1.set_token(('id', 0))

        i2 = MockInfo()
        i2.set_uri(URL('http://moth/abc.html?id=1&foo=bar'))
        i2.set_dc(QueryString([('id', ['3']),
                               ('foo', ['bar'])]))
        i2.set_token(('foo', 0))

        kb.append_uniq('a', 'b', i1)
        kb.append_uniq('a', 'b', i2)
        self.assertEqual(kb.get('a', 'b'), [i1, i2])

    def test_append_uniq_group_filter_func_specific(self):
        vuln1 = MockVuln(name='Foos')
        vuln2 = MockVuln(name='Bars')
        vuln3 = MockVuln(name='Foos', _id=42)
        kb.append_uniq_group('a', 'b', vuln1, group_klass=MockInfoSetNames)
        kb.append_uniq_group('a', 'b', vuln2, group_klass=MockInfoSetNames)
        kb.append_uniq_group('a', 'b', vuln3, group_klass=MockInfoSetNames)

        raw_data = kb.get('a', 'b')
        self.assertEqual(len(raw_data), 2)
        self.assertIsInstance(raw_data[0], InfoSet)
        self.assertIsInstance(raw_data[1], InfoSet)

        self.assertEqual(raw_data[0].get_name(), 'Foos')
        self.assertEqual(len(raw_data[0].infos), 2)
        self.assertEqual(raw_data[0].infos[1].get_id(), [42])
        self.assertEqual(raw_data[1].first_info.get_name(), 'Bars')

    def test_append_uniq_group_no_match_filter_func(self):
        vuln1 = MockVuln(name='Foos')
        vuln2 = MockVuln(name='Bars')
        kb.append_uniq_group('a', 'b', vuln1, group_klass=MockInfoSetFalse)
        info_set, created = kb.append_uniq_group('a', 'b', vuln2,
                                                 group_klass=MockInfoSetFalse)

        self.assertIsInstance(info_set, InfoSet)
        self.assertTrue(created)
        self.assertEqual(len(info_set.infos), 1)

        raw_data = kb.get('a', 'b')
        self.assertEqual(len(raw_data), 2)
        self.assertIsInstance(raw_data[0], InfoSet)
        self.assertIsInstance(raw_data[1], InfoSet)

        self.assertEqual(raw_data[0].first_info.get_name(), 'Foos')
        self.assertEqual(raw_data[1].first_info.get_name(), 'Bars')

    def test_info_set_keep_uniq_id_after_max_info_instances(self):
        #
        # Create one InfoSet, add MAX_INFO_INSTANCES, assert that the ID is not
        # changed afterwards
        #
        vuln = MockVuln(name='Foos')

        for _ in xrange(MockInfoSetNames.MAX_INFO_INSTANCES + 1):
            kb.append_uniq_group('a', 'b', vuln, group_klass=MockInfoSetNames)

        info_set_before = kb.get('a', 'b')[0]

        # Now some rounds of testing
        for _ in xrange(5):
            info_set_after, _ = kb.append_uniq_group('a', 'b', vuln,
                                                     group_klass=MockInfoSetNames)

            self.assertEqual(info_set_before.get_uniq_id(),
                             info_set_after.get_uniq_id())

    def test_append_uniq_group_filter_func_attribute_match(self):
        vuln1 = MockVuln(name='Foos')
        vuln1['tag'] = 'foo'

        vuln2 = MockVuln(name='Bars')
        vuln2['tag'] = 'bar'

        vuln3 = MockVuln(_id=42)
        vuln3['tag'] = 'foo'

        kb.append_uniq_group('a', 'b', vuln1, group_klass=MockInfoSetITag)
        kb.append_uniq_group('a', 'b', vuln2, group_klass=MockInfoSetITag)
        kb.append_uniq_group('a', 'b', vuln3, group_klass=MockInfoSetITag)

        raw_data = kb.get('a', 'b')
        self.assertEqual(len(raw_data), 2)
        self.assertIsInstance(raw_data[0], InfoSet)
        self.assertIsInstance(raw_data[1], InfoSet)

        self.assertEqual(raw_data[0].get_name(), 'Foos')
        self.assertEqual(len(raw_data[0].infos), 2)
        self.assertEqual(raw_data[0].infos[1].get_id(), [42])
        self.assertEqual(raw_data[1].first_info.get_name(), 'Bars')

    def test_strategy_run(self):
        core = w3afCore()
        
        target = core.target.get_options()
        target['target'].set_value(self.TARGET_URL)
        core.target.set_options(target)
        
        core.plugins.set_plugins(['sqli'], 'audit')
        core.plugins.init_plugins()
        
        core.verify_environment()
        core.scan_start_hook()
        
        def verify_threads_running(functor):
            thread_names = [t.name for t in threading.enumerate()]
            self.assertIn('WorkerThread', thread_names)
            self.called_teardown_audit = True
            return functor
        
        self.called_teardown_audit = False
        
        strategy = CoreStrategy(core)
        strategy._teardown_audit = verify_threads_running(strategy._teardown_audit)
        
        strategy.start()
        
        # Now test that those threads are being terminated
        self.assertTrue(self.called_teardown_audit)
        
        vulns = kb.get('sqli', 'sqli')
        self.assertEqual(len(vulns), 1, vulns)
        
        # Tell the core that we've finished, this should kill the WorkerThreads
        core.exploit_phase_prerequisites = lambda: 42
        core.scan_end_hook()

        self._assert_thread_names()