diff --git a/pyproject.toml b/pyproject.toml index e64aa0e..875843e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ dev = [ "black", "pytest-cov", "flake8", + "requests_mock" ] [tool.black] diff --git a/shbin.py b/shbin.py index a2967f5..196ab11 100644 --- a/shbin.py +++ b/shbin.py @@ -7,6 +7,7 @@ import re import secrets import sys +import requests from mimetypes import guess_extension import pyclip @@ -60,6 +61,21 @@ def get_repo_and_user(): return gh.get_repo(os.environ["SHBIN_REPO"]), gh.get_user().login +def get_repo(full_name_repo): + gh = Github(os.environ["SHBIN_GITHUB_TOKEN"]) + return gh.get_repo(full_name_repo) + + +def get_path(url): + pattern = r"https://github.com/.*?/(?:tree|blob)/(.*)" + result = re.findall(pattern, url) + if result: + extracted_string = result[0] + return extracted_string + else: + raise DocoptExit(f"Ensure your path is from github repository. (error {e})") + + def expand_paths(path_or_patterns): """ receive a list of relative paths or glob patterns and return an iterator of Path instances @@ -99,21 +115,42 @@ def download(url_or_path, repo, user): $ shbin dl https://github.com/Shiphero/pastebin/blob/main/bibo/AWS_API_fullfilment_methods/ $ shbin dl bibo/AWS_API_fullfilment_methods/ """ - path = re.sub(rf"^https://github\.com/{repo.full_name}/(blob|tree)/{repo.default_branch}/", "", url_or_path) + full_name_repo = repo.full_name + pattern = r"https://github.com/(.*?)/(tree|blob)/" + result = re.findall(pattern, url_or_path) + if result: + full_name_repo = result[0][0] + + is_from_my_repo = "https://" not in url_or_path or full_name_repo == repo.full_name + if is_from_my_repo: + path = re.sub(rf"^https://github\.com/{repo.full_name}/(blob|tree)/{repo.default_branch}/", "", url_or_path) + else: + repo = get_repo(full_name_repo) + path = get_path(url_or_path) + path = path.rstrip("/") try: - content = repo.get_contents(path) - if isinstance(content, list): - # FIXME currently this will flatten the tree: - # suposse dir/foo.py and dir/subdir/bar.py - # Then `$ shbin dl dir` will get foo.py and bar.py in the same dir. - for content_file in content: - download(content_file.path, repo, user) - return + if is_from_my_repo: + content = repo.get_contents(path) + if isinstance(content, list): + # FIXME currently this will flatten the tree: + # suposse dir/foo.py and dir/subdir/bar.py + # Then `$ shbin dl dir` will get foo.py and bar.py in the same dir. + for content_file in content: + download(content_file.path, repo, user) + return + else: + content = content.decoded_content else: - content = content.decoded_content + url = f"https://raw.githubusercontent.com/{repo.full_name}/{path}" + response = requests.get(url) + if response.status_code > 200: + raise Exception("There was a problem with your download, please check url") + content = response.content except GithubException: print("[red]x[/red] content not found") + except Exception as e: + print(f"[red]x[/red] {e}") else: target = pathlib.Path(path).name pathlib.Path(target).write_bytes(content) diff --git a/tests/test_cli.py b/tests/test_cli.py index dc5fcb6..ca40ffb 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -32,9 +32,17 @@ def repo(): repo = create_autospec(Repository, name="gh_repo") repo.create_file.return_value = {"content": Mock(html_url="https://the-url")} repo.update_file.return_value = {"content": Mock(html_url="https://the-url-updated")} + repo.full_name = "my_awesome/repository" return repo +@pytest.fixture +def outside_repo(): + outside_repo = create_autospec(Repository, name="gh_repo") + outside_repo.full_name = "another_awesome/repository" + return outside_repo + + @pytest.fixture(autouse=True) def pyclip(monkeypatch): class Stub: @@ -70,6 +78,12 @@ def patched_repo_and_user(repo): yield mocked +@pytest.fixture +def patched_another_repo_and_user(outside_repo): + with patch("shbin.get_repo", return_value=outside_repo) as mocked: + yield mocked + + @pytest.mark.parametrize("argv", (["-h"], ["--help"])) def test_help(capsys, argv): with pytest.raises(SystemExit): @@ -320,7 +334,7 @@ def test_force_new(pyclip, tmp_path, patched_repo_and_user, repo, capsys): assert capsys.readouterr().out == "🔗📋 https://the-url-2\n" -def test_download_a_file(tmp_path, patched_repo_and_user, repo): +def test_download_a_file_from_owned_repo(tmp_path, patched_repo_and_user, repo): git_data = { "decoded_content": b"awesome content", } @@ -330,3 +344,29 @@ def test_download_a_file(tmp_path, patched_repo_and_user, repo): os.chdir(working_dir) main(["dl", "hello.md"]) assert (working_dir / "hello.md").read_bytes() == b"awesome content" + + +def test_download_a_file_from_public_repo( + tmp_path, patched_another_repo_and_user, outside_repo, requests_mock, patched_repo_and_user +): + requests_mock.get( + "https://raw.githubusercontent.com/another_awesome/repository/main/hello.md", content=b"awesome content" + ) + working_dir = tmp_path / "working_dir" + working_dir.mkdir() + os.chdir(working_dir) + main(["dl", "https://github.com/another_awesome/repository/blob/main/hello.md"]) + assert (working_dir / "hello.md").read_bytes() == b"awesome content" + + +def test_download_a_file_from_public_repo_rasie_error( + tmp_path, patched_another_repo_and_user, outside_repo, requests_mock, patched_repo_and_user +): + requests_mock.get("https://raw.githubusercontent.com/another_awesome/repository/main/hello.md", status_code=400) + working_dir = tmp_path / "working_dir" + working_dir.mkdir() + os.chdir(working_dir) + main(["dl", "https://github.com/another_awesome/repository/blob/main/hello.md"]) + with pytest.raises(Exception) as exc_info: + raise Exception("There was a problem with your download, please check url") + assert str(exc_info.value) == "There was a problem with your download, please check url"