async def test_set_proxy_set_aiohttp_engine_proxy(self, loop): h = HammerTime(loop=loop, request_engine=MagicMock()) h.set_proxy("http://some.proxy.com/") aiohttp_engine = h.request_engine.request_engine aiohttp_engine.set_proxy.assert_called_with("http://some.proxy.com/")
async def test_interrupt_close_hammertime(self, loop): hammertime = HammerTime(loop=loop) hammertime._interrupt() # Wait for hammertime.close to be called. await hammertime.closed self.assertTrue(hammertime.is_closed)
async def test_constructor_do_not_overwrite_aiohttp_engine_proxy_if_constructor_proxy_is_none( self, loop): engine = AioHttpEngine(loop=loop, proxy="http://some.proxy.com") h = HammerTime(loop=loop, request_engine=engine) aiohttp_engine = h.request_engine.request_engine self.assertEqual(aiohttp_engine.proxy, "http://some.proxy.com")
async def test_request_raise_cancelled_error_if_hammertime_is_close( self, loop): hammertime = HammerTime(loop=loop) await hammertime.close() self.assertTrue(hammertime.is_closed) with self.assertRaises(asyncio.CancelledError): await hammertime.request("http://example.com")
async def test_explicit_abandon_obtained_when_requested(self, loop): h = HammerTime(loop=loop, request_engine=FakeEngine(), retry_count=2) h.heuristics.add(RejectStatusCode(range(0, 600))) # Everything future = h.request("http://example.com/1") with self.assertRaises(RejectRequest): await future self.assertEqual(0, h.stats.retries)
async def test_constructor_set_aiohttp_engine_proxy_if_constructor_proxy_is_not_none( self, loop): h = HammerTime(loop=loop, request_engine=MagicMock(), proxy="http://some.proxy.com/") aiohttp_engine = h.request_engine.request_engine aiohttp_engine.set_proxy.assert_called_once_with( "http://some.proxy.com/")
async def test_fetch_session_cookies_on_scan_start_if_no_user_supplied_cookies( self, loop): engine = FakeHammerTimeEngine() hammertime = HammerTime(request_engine=engine, loop=loop) hammertime.collect_successful_requests() await tachyon.scan(hammertime, cookies=None, accumulator=self.accumulator) tachyon.get_session_cookies.assert_called_once_with(hammertime)
async def test_provide_exception_when_resolving_specific_promise( self, loop): h = HammerTime(loop=loop, request_engine=FakeEngine()) h.heuristics.add(BlockRequest("http://example.com/1")) future = h.request("http://example.com/1") with self.assertRaises(StopRequest): await future with self.assertRaises(StopRequest): await future
async def test_successful_requests_return_if_no_pending_requests( self, loop): h = HammerTime(loop=loop, request_engine=FakeEngine()) h.collect_successful_requests() try: with async_timeout.timeout(0.001): async for entry in h.successful_requests(): entry.request except asyncio.TimeoutError: self.fail("Function blocked.")
async def test_paths_exists_fetch_generated_paths(self, loop): path_generator = MagicMock() path_generator.generate_paths.return_value = ["/", "/test", "/path"] fake_directory_fetcher = MagicMock() fake_directory_fetcher.fetch_paths = make_mocked_coro() tachyon.PathGenerator = MagicMock(return_value=path_generator) tachyon.DirectoryFetcher = MagicMock(return_value=fake_directory_fetcher) await tachyon.test_paths_exists(HammerTime(loop=loop)) fake_directory_fetcher.fetch_paths.assert_called_once_with(path_generator.generate_paths.return_value)
async def test_loop_over_results(self, loop): h = HammerTime(loop=loop, request_engine=FakeEngine()) h.request("http://example.com/1") h.request("http://example.com/2") out = set() async for entry in h.successful_requests(): out.add(entry.response.content) self.assertEqual(out, {"http://example.com/1", "http://example.com/2"}) self.assertEqual(h.completed_count, 2)
async def test_scan_plugins_only(self, loop): engine = FakeHammerTimeEngine() hammertime = HammerTime(request_engine=engine, loop=loop) hammertime.collect_successful_requests() await tachyon.scan(hammertime, plugins_only=True, accumulator=self.accumulator) tachyon.load_execute_host_plugins.assert_called_once_with(hammertime) tachyon.test_file_exists.assert_not_called() tachyon.test_paths_exists.assert_not_called()
async def test_file_exists_fetch_all_generate_files(self, loop): database.valid_paths = ["/path/file%d" % i for i in range(10)] fake_file_fetcher = MagicMock() fake_file_fetcher.fetch_files = make_mocked_coro() tachyon.FileFetcher = MagicMock(return_value=fake_file_fetcher) fake_file_generator = MagicMock() fake_file_generator.generate_files.return_value = ["list of files"] with patch("tachyon.__main__.FileGenerator", MagicMock(return_value=fake_file_generator)): await tachyon.test_file_exists(HammerTime(loop=loop)) fake_file_fetcher.fetch_files.assert_called_once_with(["list of files"])
async def test_wait_for_multiple_requests(self, loop): h = HammerTime(loop=loop, request_engine=FakeEngine()) promise_1 = h.request("http://example.com/1") promise_2 = h.request("http://example.com/2") entry = await promise_2 self.assertEqual(entry.response.content, "http://example.com/2") entry = await promise_1 self.assertEqual(entry.response.content, "http://example.com/1") await h.close() self.assertEqual(h.completed_count, 2)
async def test_paths_exists_output_fetch_paths_count(self, loop): path_generator = MagicMock() paths = ["/", "/test", "/path"] path_generator.generate_paths.return_value = paths fake_directory_fetcher = MagicMock() fake_directory_fetcher.fetch_paths = make_mocked_coro() tachyon.PathGenerator = MagicMock(return_value=path_generator) tachyon.DirectoryFetcher = MagicMock(return_value=fake_directory_fetcher) with patch("tachyon.textutils.output_info") as output_info: await tachyon.test_paths_exists(HammerTime(loop=loop)) output_info.assert_any_call("Probing %d paths" % len(paths))
async def test_skip_results_that_fail(self, loop): h = HammerTime(loop=loop, request_engine=FakeEngine()) h.heuristics.add(BlockRequest("http://example.com/1")) h.request("http://example.com/1") h.request("http://example.com/2") out = set() async for entry in h.successful_requests(): out.add(entry.response.content) self.assertEqual(out, {"http://example.com/2"}) self.assertEqual(h.completed_count, 2)
async def test_enumerate_found_raise_if_host_is_offline(self, loop): async def fake_perform(*args, **kwargs): raise OfflineHostException() hammertime = HammerTime(loop) hammertime.request_engine = MagicMock() hammertime.request_engine.perform = fake_perform component_finder = ActiveComponentFinder(hammertime, self.target_url) component_finder.components_file_list_group = self.plugin_list with self.assertRaises(OfflineHostException): async for _ in component_finder.enumerate_found(): pass
async def test_use_user_supplied_cookies_if_available(self, loop): database.session_cookie = "my-cookies=123" cookies = "test-cookie=true" engine = FakeHammerTimeEngine() hammertime = HammerTime(request_engine=engine, loop=loop) hammertime.collect_successful_requests() with patch("tachyon.config.add_http_header") as add_http_header: await tachyon.scan(hammertime, cookies=cookies, accumulator=self.accumulator) add_http_header.assert_any_call(ANY, "Cookie", "test-cookie=true")
async def test_no_successful_request_returned_when_requests_are_cancelled( self, loop): engine = MagicMock() engine.perform = make_mocked_coro( raise_exception=asyncio.CancelledError) hammertime = HammerTime(loop=loop, request_engine=engine) for i in range(5): hammertime.request("http://example.com") successful_requests = [] async for request in hammertime.successful_requests(): successful_requests.append(request) self.assertEqual(successful_requests, [])
async def test_paths_exists_output_paths_found_count(self, loop): path_generator = MagicMock() paths = ["/", "/test", "/path"] path_generator.generate_paths.return_value = paths fake_directory_fetcher = MagicMock() fake_directory_fetcher.fetch_paths = make_mocked_coro() tachyon.PathGenerator = MagicMock(return_value=path_generator) tachyon.DirectoryFetcher = MagicMock(return_value=fake_directory_fetcher) asyncio.set_event_loop(asyncio.new_event_loop()) database.valid_paths = paths with patch("tachyon.textutils.output_info") as output_info: await tachyon.test_paths_exists(HammerTime(loop=loop)) output_info.assert_any_call("Found %d valid paths" % len(database.valid_paths))
async def test_paths_exists_do_recursive_path_search_if_recursive_is_true(self, loop): path_generator = MagicMock() paths = ["/", "/test", "/path"] path_generator.generate_paths.return_value = paths fake_directory_fetcher = MagicMock() fake_directory_fetcher.fetch_paths = make_mocked_coro() tachyon.PathGenerator = MagicMock(return_value=path_generator) tachyon.DirectoryFetcher = MagicMock(return_value=fake_directory_fetcher) await tachyon.test_paths_exists(HammerTime(loop=loop), recursive=True) path_generator.generate_paths.assert_has_calls([call(use_valid_paths=False), call(use_valid_paths=True), call(use_valid_paths=True)], any_order=False) fake_directory_fetcher.fetch_paths.assert_has_calls([call(paths)]*3)
async def test_enumerate_found_return_list_of_dict_with_component_key_and_fetched_files( self, loop): async def fake_perform(entry, *args, **kwargs): entry.result.hash = "fake-hash" return entry hammertime = HammerTime(loop) hammertime.request_engine = MagicMock() hammertime.request_engine.perform = fake_perform component_finder = ActiveComponentFinder(hammertime, self.target_url) component_finder.components_file_list_group = self.plugin_list components = await self.return_async_iterator_as_list( component_finder.enumerate_found()) for component_dict in components: self.assertTrue( component_dict['key'] == self.plugin0_file_list.key or component_dict['key'] == self.plugin1_file_list.key) self.assertIn('files', component_dict)
async def test_add_http_headers(self, loop): hammertime = HammerTime(loop=loop) config.heuristics_with_child = [ RejectCatchAllRedirect(), FollowRedirects() ] hammertime.heuristics.add_multiple(config.heuristics_with_child) hammertime.heuristics.add = MagicMock() for heuristic in config.heuristics_with_child: heuristic.child_heuristics.add = MagicMock() config.add_http_header(hammertime, "header", "value") set_header = hammertime.heuristics.add.call_args[0][0] self.assertEqual(set_header.name, "header") self.assertEqual(set_header.value, "value") for heuristic_with_child in config.heuristics_with_child: set_header = heuristic_with_child.child_heuristics.add.call_args[ 0][0] self.assertEqual(set_header.name, "header") self.assertEqual(set_header.value, "value")
async def test_add_requests_and_wait_for_completion(self, loop): h = HammerTime(loop=loop, request_engine=FakeEngine()) entry = await h.request("http://example.com") self.assertEqual(entry.response.content, "http://example.com") await h.close()
async def test_retries_performed_and_response_obtained(self, loop): h = HammerTime(loop=loop, request_engine=FakeEngine(), retry_count=2) h.heuristics.add(BlockRequest("http://example.com/1")) entry = await h.request("http://example.com/1") self.assertEqual(entry.response.content, "http://example.com/1")
def async_setup(self, loop): self.hammertime = HammerTime(loop=loop, request_engine=FakeHammerTimeEngine()) self.hammertime.collect_successful_requests() self.hammertime.heuristics.add_multiple([SetFlagInResult("soft404", False), SetFlagInResult("error_behavior", False)]) self.directory_fetcher = DirectoryFetcher(self.host, self.hammertime)
async def test_open_and_close(self, loop): h = HammerTime(loop=loop) await h.close() self.assertEqual(h.completed_count, 0)
def _setup_async_test(self, loop): self.hammertime = HammerTime(loop=loop) self.fetcher = FileFetcher(self.hammertime, "www.example.com")
async def test_preserve_arguments(self, loop): h = HammerTime(loop=loop, request_engine=FakeEngine()) entry = await h.request("http://example.com", arguments={"hello": "world"}) self.assertEqual(entry.arguments["hello"], "world")