diff --git a/tests/discovery/test_githubcode.py b/tests/discovery/test_githubcode.py index 275b9100..bf0b7493 100644 --- a/tests/discovery/test_githubcode.py +++ b/tests/discovery/test_githubcode.py @@ -39,6 +39,31 @@ class TestSearchGithubCode: test_result = test_class_instance.fragments_from_response(response) assert test_result == ["test1", "test2"] + def test_fail_fragments_from_response(self): + Core.github_key = MagicMock(return_value="lol") + test_class_instance = githubcode.SearchGithubCode(word="test", limit=500) + response = Response() + json = { + "items": [ + { + "fail": True + }, + { + "text_matches": [] + }, + { + "text_matches": [ + { + "weird": "result" + } + ] + } + ] + } + response.json = MagicMock(return_value=json) + test_result = test_class_instance.fragments_from_response(response) + assert test_result == [] + if __name__ == '__main__': pytest.main() diff --git a/theHarvester/discovery/githubcode.py b/theHarvester/discovery/githubcode.py index 0fb20664..0dbb5625 100644 --- a/theHarvester/discovery/githubcode.py +++ b/theHarvester/discovery/githubcode.py @@ -42,44 +42,31 @@ class SearchGithubCode: @staticmethod def fragments_from_response(response: Response) -> List[str]: - items: List[Dict[str, Any]] = response.json().get('items') - fragments = [] - if items is not None: - for item in items: - matches = item.get("text_matches") - for match in matches: - fragments.append(match.get("fragment")) - return fragments + items: List[Dict[str, Any]] = response.json().get('items') or list() + fragments: List[str] = list() + for item in items: + matches = item.get("text_matches") or list() + for match in matches: + fragments.append(match.get("fragment")) + return [fragment for fragment in fragments if fragment is not None] @staticmethod - def next_page_from_response(response: Response) -> Optional[int]: - next_link = response.links.get("next") - if next_link: - parsed = urlparse.urlparse(next_link.get("url")) + def page_from_response(page: str, response: Response) -> Optional[int]: + page_link = response.links.get(page) + if page_link: + parsed = urlparse.urlparse(page_link.get("url")) params = urlparse.parse_qs(parsed.query) page = params.get('page') or [None] - next_page = page[0] and int(page[0]) - return next_page - else: - return None - - @staticmethod - def last_page_from_response(response: Response) -> Optional[int]: - next_link = response.links.get("last") - if next_link: - parsed = urlparse.urlparse(next_link.get("url")) - params = urlparse.parse_qs(parsed.query) - page = params.get('page') or [None] - last_page = page[0] and int(page[0]) - return last_page + page_number = page[0] and int(page[0]) + return page_number else: return None def handle_response(self, response: Response) -> Optional[Any]: if response.ok: fragments = self.fragments_from_response(response) - next_page = self.next_page_from_response(response) - last_page = self.last_page_from_response(response) + next_page = self.page_from_response("next", response) + last_page = self.page_from_response("last", response) return SuccessResult(fragments, next_page, last_page) elif response.status_code == 429 or response.status_code == 403: return RetryResult(60)