From a4b846e58c3dc5b28119189a6c8fc0d093c19a23 Mon Sep 17 00:00:00 2001 From: Giovanni M Guidini <99758426+giovanni-guidini@users.noreply.github.com> Date: Tue, 5 Sep 2023 18:11:57 -0300 Subject: [PATCH] feat: support gh refresh tokens (#69) Depends on codecov/shared#27 Adds support for github app refresh tokens context: codecov/engineering-team#162 --- conftest.py | 6 ++-- helpers/token_refresh.py | 35 ++++++++++++++++++++++ requirements.in | 2 +- requirements.txt | 2 +- services/owner.py | 7 ++++- services/repository.py | 32 ++------------------ services/tests/test_repository_service.py | 36 +++++++++++++++++++++-- 7 files changed, 81 insertions(+), 39 deletions(-) create mode 100644 helpers/token_refresh.py diff --git a/conftest.py b/conftest.py index 734f46eed..54ca45107 100644 --- a/conftest.py +++ b/conftest.py @@ -271,12 +271,10 @@ def mock_submit_fn(metric, start, end): mock_submit = mocker.Mock() mock_submit.side_effect = mock_submit_fn - import helpers + from helpers.checkpoint_logger import CheckpointLogger return mocker.patch.object( - helpers.checkpoint_logger.CheckpointLogger, + CheckpointLogger, "submit_subflow", mock_submit, ) - - return mock_submit diff --git a/helpers/token_refresh.py b/helpers/token_refresh.py new file mode 100644 index 000000000..dcb380f72 --- /dev/null +++ b/helpers/token_refresh.py @@ -0,0 +1,35 @@ +import logging +from typing import Callable, Dict + +from shared.encryption.token import encode_token + +from database.models.core import Owner +from services.encryption import encryptor + +log = logging.getLogger(__name__) + + +def get_token_refresh_callback(owner: Owner) -> Callable[[Dict], None]: + """ + Produces a callback function that will encode and update the oauth token of a user. + This callback is passed to the TorngitAdapter for the service. + """ + # Some tokens don't have to be refreshed (GH integration, default bots) + # They don't belong to any owners. + if owner is None: + return None + + service = owner.service + if service == "bitbucket" or service == "bitbucket_server": + return None + + async def callback(new_token: Dict) -> None: + log.info( + "Saving new token after refresh", + extra=dict(owner=owner.username, ownerid=owner.ownerid), + ) + string_to_save = encode_token(new_token) + oauth_token = encryptor.encode(string_to_save).decode() + owner.oauth_token = oauth_token + + return callback diff --git a/requirements.in b/requirements.in index 496cb182c..ccbee6ceb 100644 --- a/requirements.in +++ b/requirements.in @@ -1,4 +1,4 @@ -git+ssh://git@github.com/codecov/shared.git@6a8a33248804a9c101d34f417efda7c11e4bbe63#egg=shared +git+ssh://git@github.com/codecov/shared.git@f0174635ccaebc3463a5641b9ad640a46b5fd472#egg=shared git+ssh://git@github.com/codecov/opentelem-python.git@v0.0.4a1#egg=codecovopentelem boto3 celery diff --git a/requirements.txt b/requirements.txt index 32d6df2e4..cac8f38ed 100644 --- a/requirements.txt +++ b/requirements.txt @@ -251,7 +251,7 @@ s3transfer==0.3.4 # via boto3 sentry-sdk==1.19.1 # via -r requirements.in -shared @ git+ssh://git@github.com/codecov/shared.git@4f65b0040ab6bb3dba2190ce544ade9a0ea3e354 +shared @ git+ssh://git@github.com/codecov/shared.git@f0174635ccaebc3463a5641b9ad640a46b5fd472 # via -r requirements.in six==1.15.0 # via diff --git a/services/owner.py b/services/owner.py index 6a71e318e..5f57bafae 100644 --- a/services/owner.py +++ b/services/owner.py @@ -3,7 +3,7 @@ import shared.torngit as torngit from shared.config import get_config, get_verify_ssl -from database.models import Owner +from helpers.token_refresh import get_token_refresh_callback from services.bots import get_owner_appropriate_bot_token log = logging.getLogger(__name__) @@ -27,6 +27,11 @@ def get_owner_provider_service(owner, using_integration=False): key=get_config(service, "client_id"), secret=get_config(service, "client_secret"), ), + # if using integration we will use the integration token + # not the owner's token + on_token_refresh=( + get_token_refresh_callback(owner) if not using_integration else None + ), ) return _get_owner_provider_service_instance(service, **adapter_params) diff --git a/services/repository.py b/services/repository.py index 41d4dad8c..c4598aa28 100644 --- a/services/repository.py +++ b/services/repository.py @@ -2,11 +2,10 @@ import re from dataclasses import dataclass from datetime import datetime -from typing import Any, Callable, Dict, Mapping, Optional, Tuple +from typing import Any, Mapping, Optional, Tuple import shared.torngit as torngit from shared.config import get_config, get_verify_ssl -from shared.encryption.token import encode_token from shared.torngit.exceptions import ( TorngitClientError, TorngitError, @@ -14,11 +13,10 @@ ) from sqlalchemy.dialects.postgresql import insert from sqlalchemy.exc import IntegrityError -from sqlalchemy.orm import Session from database.models import Commit, Owner, Pull, Repository +from helpers.token_refresh import get_token_refresh_callback from services.bots import get_repo_appropriate_bot_token, get_token_type_mapping -from services.encryption import encryptor from services.yaml import read_yaml_field log = logging.getLogger(__name__) @@ -26,32 +24,6 @@ merged_pull = re.compile(r".*Merged in [^\s]+ \(pull request \#(\d+)\).*").match -def get_token_refresh_callback(owner: Owner) -> Callable[[Dict], None]: - """ - Produces a callback function that will encode and update the oauth token of a user. - This callback is passed to the TorngitAdapter for the service. - """ - # Some tokens don't have to be refreshed (GH integration, default bots) - # They don't belong to any owners. - if owner is None: - return None - - service = owner.service - if service != "gitlab" and service != "gitlab_enterprise": - return None - - async def callback(new_token: Dict) -> None: - log.info( - "Saving new token after refresh", - extra=dict(owner=owner.username, ownerid=owner.ownerid), - ) - string_to_save = encode_token(new_token) - oauth_token = encryptor.encode(string_to_save).decode() - owner.oauth_token = oauth_token - - return callback - - def get_repo_provider_service( repository, commit=None ) -> torngit.base.TorngitBaseAdapter: diff --git a/services/tests/test_repository_service.py b/services/tests/test_repository_service.py index 1d43ffb7e..481769ebe 100644 --- a/services/tests/test_repository_service.py +++ b/services/tests/test_repository_service.py @@ -31,7 +31,7 @@ class TestRepositoryServiceTestCase(object): - def test_get_repo_provider_service(self, dbsession): + def test_get_repo_provider_service_github(self, dbsession): repo = RepositoryFactory.create( owner__unencrypted_oauth_token="testyftq3ovzkb3zmt823u3t04lkrt9w", owner__service="github", @@ -55,7 +55,39 @@ def test_get_repo_provider_service(self, dbsession): } assert res.data == expected_data assert repo.owner.service == "github" - assert res._on_token_refresh is None # GH doesn't have callback implemented + assert res._on_token_refresh is not None + assert inspect.isawaitable(res._on_token_refresh(None)) + assert res.token == { + "username": repo.owner.username, + "key": "testyftq3ovzkb3zmt823u3t04lkrt9w", + "secret": None, + } + + def test_get_repo_provider_service_bitbucket(self, dbsession): + repo = RepositoryFactory.create( + owner__unencrypted_oauth_token="testyftq3ovzkb3zmt823u3t04lkrt9w", + owner__service="bitbucket", + name="example-python", + ) + dbsession.add(repo) + dbsession.flush() + res = get_repo_provider_service(repo) + expected_data = { + "owner": { + "ownerid": repo.owner.ownerid, + "service_id": repo.owner.service_id, + "username": repo.owner.username, + }, + "repo": { + "name": "example-python", + "using_integration": False, + "service_id": repo.service_id, + "repoid": repo.repoid, + }, + } + assert res.data == expected_data + assert repo.owner.service == "bitbucket" + assert res._on_token_refresh is None assert res.token == { "username": repo.owner.username, "key": "testyftq3ovzkb3zmt823u3t04lkrt9w",