diff --git a/tests/discovery/test_githubcode_additions.py b/tests/discovery/test_githubcode_additions.py new file mode 100644 index 00000000..d156dde8 --- /dev/null +++ b/tests/discovery/test_githubcode_additions.py @@ -0,0 +1,79 @@ +from unittest.mock import MagicMock, AsyncMock +import asyncio +import pytest +from _pytest.mark.structures import MarkDecorator +from theHarvester.discovery import githubcode +from theHarvester.lib.core import Core + +pytestmark: MarkDecorator = pytest.mark.asyncio + + +class TestSearchGithubCodeProcess: + async def test_process_stops_after_max_retries(self, monkeypatch): + Core.github_key = MagicMock(return_value="test_key") # type: ignore[method-assign] + inst = githubcode.SearchGithubCode(word="test", limit=10) + + # Speed up by avoiding actual sleeps + monkeypatch.setattr(githubcode, "get_delay", lambda: 0, raising=False) + monkeypatch.setattr(asyncio, "sleep", AsyncMock(return_value=None)) + + # Force RetryResult every time + monkeypatch.setattr( + inst, + "handle_response", + AsyncMock(return_value=githubcode.RetryResult(0)), + ) + monkeypatch.setattr( + inst, + "do_search", + AsyncMock(return_value=("", {}, 403, {})), + ) + + inst.max_retries = 2 + await inst.process() + assert inst.page == 0, "Process should stop after exceeding max retries" + assert inst.retry_count == 3, "Retry count should exceed max_retries before stopping" + + async def test_process_stops_on_error_result(self, monkeypatch): + Core.github_key = MagicMock(return_value="test_key") # type: ignore[method-assign] + inst = githubcode.SearchGithubCode(word="test", limit=10) + + monkeypatch.setattr(githubcode, "get_delay", lambda: 0, raising=False) + monkeypatch.setattr(asyncio, "sleep", AsyncMock(return_value=None)) + + # Force ErrorResult + monkeypatch.setattr( + inst, + "handle_response", + AsyncMock(return_value=githubcode.ErrorResult(500, "err")), + ) + monkeypatch.setattr( + inst, + "do_search", + AsyncMock(return_value=("", {}, 500, {})), + ) + + await inst.process() + assert inst.page == 0, "Process should stop on error result to avoid infinite loop" + + async def test_process_breaks_on_same_page_pagination(self, monkeypatch): + Core.github_key = MagicMock(return_value="test_key") # type: ignore[method-assign] + inst = githubcode.SearchGithubCode(word="test", limit=10) + + monkeypatch.setattr(githubcode, "get_delay", lambda: 0, raising=False) + monkeypatch.setattr(asyncio, "sleep", AsyncMock(return_value=None)) + + # Force SuccessResult that does not advance the page + monkeypatch.setattr( + inst, + "handle_response", + AsyncMock(return_value=githubcode.SuccessResult([], next_page=1, last_page=0)), + ) + monkeypatch.setattr( + inst, + "do_search", + AsyncMock(return_value=("", {"items": []}, 200, {})), + ) + + await inst.process() + assert inst.page == 0, "Process should stop when pagination does not advance" diff --git a/theHarvester/discovery/githubcode.py b/theHarvester/discovery/githubcode.py index 8edb7c5f..adbf487d 100644 --- a/theHarvester/discovery/githubcode.py +++ b/theHarvester/discovery/githubcode.py @@ -45,6 +45,9 @@ class SearchGithubCode: 'Accept': 'application/vnd.github.v3.text-match+json', 'Authorization': f'token {self.key}', } + # Retry control to avoid infinite loops on rate limiting + self.retry_count = 0 + self.max_retries = 3 except Exception as e: print(f'Error initializing SearchGithubCode: {e}') raise @@ -116,17 +119,33 @@ class SearchGithubCode: result = await self.handle_response(api_response) if isinstance(result, SuccessResult): + # Reset retry counter on any successful response + self.retry_count = 0 print(f'\tSearching {self.counter} results.') self.total_results += ''.join(result.fragments) self.counter += len(result.fragments) - self.page = result.next_page or result.last_page + next_or_last = result.next_page or result.last_page + # Break if pagination does not advance to avoid infinite loop + if next_or_last == self.page: + print('\tNo page advancement detected; exiting to avoid infinite loop.') + self.page = 0 + break + self.page = next_or_last await asyncio.sleep(get_delay()) elif isinstance(result, RetryResult): + self.retry_count += 1 + if self.retry_count > self.max_retries: + print('\tMaximum retries reached; exiting to avoid infinite loop.') + self.page = 0 + break sleepy_time = get_delay() + result.time print(f'\tRetrying page in {sleepy_time} seconds...') await asyncio.sleep(sleepy_time) else: + # On error, stop to avoid endless retries on a bad state print(f'\tException occurred: status_code: {result.status_code} reason: {result.body}') + self.page = 0 + break except Exception as e: print(f'Error processing page: {e}') await asyncio.sleep(get_delay())