<commit_before>import time
from functools import wraps

from itsdangerous import JSONWebSignatureSerializer
from flask import Blueprint, request, abort, jsonify, current_app

from flask_oauthlib.provider import OAuth2Provider
from flask_oauthlib.provider import OAuth2RequestValidator
from flask_oauthlib.provider.oauth2 import log
from flask_oauthlib.utils import decode_base64
from oauthlib.common import to_unicode
from oauthlib.oauth2.rfc6749.tokens import random_token_generator

from skylines.database import db
from skylines.model import User, Client, RefreshToken, AccessToken
from skylines.lib.types import is_bytes


class CustomProvider(OAuth2Provider):
    def __init__(self, *args, **kwargs):
        super(CustomProvider, self).__init__(*args, **kwargs)
        self.blueprint = Blueprint("oauth", __name__)

    def init_app(self, app):
        super(CustomProvider, self).init_app(app)
        app.config.setdefault("OAUTH2_PROVIDER_TOKEN_GENERATOR", self.generate_token)
        app.config.setdefault(
            "OAUTH2_PROVIDER_REFRESH_TOKEN_GENERATOR", random_token_generator
        )

        app.jws = JSONWebSignatureSerializer(app.config.get("SECRET_KEY"))

        app.register_blueprint(self.blueprint)

    def generate_token(self, request):
        token = {"user": request.user.id, "exp": int(time.time() + request.expires_in)}

        if request.scopes is not None:
            token["scope"] = " ".join(request.scopes)

        return current_app.jws.dumps(token).decode("ascii")

    def verify_request(self, scopes):
        if request.authorization:
            from skylines.model import User

            username = request.authorization.username
            if is_bytes(username):
                username = username.decode("utf-8")

            password = request.authorization.password
            if is_bytes(password):
                password = password.decode("utf-8")

            user = User.by_credentials(username, password)

            request.user_id = user.id if user else None
            return (user is not None), None

        else:
            valid, req = super(CustomProvider, self).verify_request(scopes)

            request.user_id = req.access_token.user_id if valid else None

            return valid, req

    def required(self, *args, **kwargs):
        return self.require_oauth(*args, **kwargs)

    def optional(self, *scopes):
        """Enhance resource with specified scopes."""

        def wrapper(f):
            @wraps(f)
            def decorated(*args, **kwargs):
                for func in self._before_request_funcs:
                    func()

                if hasattr(request, "oauth") and request.oauth:
                    return f(*args, **kwargs)

                valid, req = self.verify_request(scopes)

                for func in self._after_request_funcs:
                    valid, req = func(valid, req)

                if not valid and (
                    not req or "Authorization" in req.headers or req.access_token
                ):
                    if self._invalid_response:
                        return self._invalid_response(req)
                    return abort(401)
                request.oauth = req
                return f(*args, **kwargs)

            return decorated

        return wrapper


class CustomRequestValidator(OAuth2RequestValidator):
    def __init__(self):
        super(CustomRequestValidator, self).__init__(
            clientgetter=lambda client_id: Client(),
            tokengetter=self.tokengetter,
            grantgetter=None,
            usergetter=User.by_credentials,
            tokensetter=self.tokensetter,
        )

    @staticmethod
    def tokengetter(access_token=None, refresh_token=None):
        """ Retrieve a token record using submitted access token or refresh token. """
        if access_token:
            return AccessToken.from_jwt(access_token)

        elif refresh_token:
            return RefreshToken.query(refresh_token=refresh_token).first()

    @staticmethod
    def tokensetter(token, request, *args, **kwargs):
        """Save a new token to the database.

        :param token: Token dictionary containing access and refresh tokens, plus token type.
        :param request: Request dictionary containing information about the client and user.
        """

        if request.grant_type != "refresh_token":
            tok = RefreshToken(
                refresh_token=token["refresh_token"], user_id=request.user.id
            )
            db.session.add(tok)
            db.session.commit()

    def rotate_refresh_token(self, request):
        return False

    def authenticate_client(self, request, *args, **kwargs):

        auth = request.headers.get("Authorization", None)
        if auth:
            try:
                _, s = auth.split(" ")
                client_id, client_secret = decode_base64(s).split(":")
                client_id = to_unicode(client_id, "utf-8")
            except Exception as e:
                log.debug("Authenticate client failed with exception: %r", e)
                return False
        else:
            client_id = request.client_id

        client = self._clientgetter(client_id)
        if not client:
            log.debug("Authenticate client failed, client not found.")
            return False

        return self.authenticate_client_id(client_id, request)


oauth = CustomProvider()
oauth._validator = CustomRequestValidator()


@oauth.blueprint.route("/oauth/token", methods=["POST"])
@oauth.token_handler
def access_token(*args, **kwargs):
    return None


@oauth.blueprint.route("/oauth/revoke", methods=["POST"])
@oauth.revoke_handler
def revoke_token():
    pass


@oauth.invalid_response
def invalid_require_oauth(req):
    message = req.error_message if req else "Unauthorized"
    return jsonify(error="invalid_token", message=message), 401
<commit_msg>api/oauth: Add `user_id` to Sentry events<commit_after>import time
from functools import wraps

from itsdangerous import JSONWebSignatureSerializer
from flask import Blueprint, request, abort, jsonify, current_app

from flask_oauthlib.provider import OAuth2Provider
from flask_oauthlib.provider import OAuth2RequestValidator
from flask_oauthlib.provider.oauth2 import log
from flask_oauthlib.utils import decode_base64
from oauthlib.common import to_unicode
from oauthlib.oauth2.rfc6749.tokens import random_token_generator
from sentry_sdk import set_user as set_sentry_user

from skylines.database import db
from skylines.model import User, Client, RefreshToken, AccessToken
from skylines.lib.types import is_bytes


class CustomProvider(OAuth2Provider):
    def __init__(self, *args, **kwargs):
        super(CustomProvider, self).__init__(*args, **kwargs)
        self.blueprint = Blueprint("oauth", __name__)

    def init_app(self, app):
        super(CustomProvider, self).init_app(app)
        app.config.setdefault("OAUTH2_PROVIDER_TOKEN_GENERATOR", self.generate_token)
        app.config.setdefault(
            "OAUTH2_PROVIDER_REFRESH_TOKEN_GENERATOR", random_token_generator
        )

        app.jws = JSONWebSignatureSerializer(app.config.get("SECRET_KEY"))

        app.register_blueprint(self.blueprint)

    def generate_token(self, request):
        token = {"user": request.user.id, "exp": int(time.time() + request.expires_in)}

        if request.scopes is not None:
            token["scope"] = " ".join(request.scopes)

        return current_app.jws.dumps(token).decode("ascii")

    def verify_request(self, scopes):
        if request.authorization:
            from skylines.model import User

            username = request.authorization.username
            if is_bytes(username):
                username = username.decode("utf-8")

            password = request.authorization.password
            if is_bytes(password):
                password = password.decode("utf-8")

            user = User.by_credentials(username, password)

            request.user_id = user.id if user else None
            set_sentry_user({"id": request.user_id})

            return (user is not None), None

        else:
            valid, req = super(CustomProvider, self).verify_request(scopes)

            request.user_id = req.access_token.user_id if valid else None
            set_sentry_user({"id": request.user_id})

            return valid, req

    def required(self, *args, **kwargs):
        return self.require_oauth(*args, **kwargs)

    def optional(self, *scopes):
        """Enhance resource with specified scopes."""

        def wrapper(f):
            @wraps(f)
            def decorated(*args, **kwargs):
                for func in self._before_request_funcs:
                    func()

                if hasattr(request, "oauth") and request.oauth:
                    return f(*args, **kwargs)

                valid, req = self.verify_request(scopes)

                for func in self._after_request_funcs:
                    valid, req = func(valid, req)

                if not valid and (
                    not req or "Authorization" in req.headers or req.access_token
                ):
                    if self._invalid_response:
                        return self._invalid_response(req)
                    return abort(401)
                request.oauth = req
                return f(*args, **kwargs)

            return decorated

        return wrapper


class CustomRequestValidator(OAuth2RequestValidator):
    def __init__(self):
        super(CustomRequestValidator, self).__init__(
            clientgetter=lambda client_id: Client(),
            tokengetter=self.tokengetter,
            grantgetter=None,
            usergetter=User.by_credentials,
            tokensetter=self.tokensetter,
        )

    @staticmethod
    def tokengetter(access_token=None, refresh_token=None):
        """ Retrieve a token record using submitted access token or refresh token. """
        if access_token:
            return AccessToken.from_jwt(access_token)

        elif refresh_token:
            return RefreshToken.query(refresh_token=refresh_token).first()

    @staticmethod
    def tokensetter(token, request, *args, **kwargs):
        """Save a new token to the database.

        :param token: Token dictionary containing access and refresh tokens, plus token type.
        :param request: Request dictionary containing information about the client and user.
        """

        if request.grant_type != "refresh_token":
            tok = RefreshToken(
                refresh_token=token["refresh_token"], user_id=request.user.id
            )
            db.session.add(tok)
            db.session.commit()

    def rotate_refresh_token(self, request):
        return False

    def authenticate_client(self, request, *args, **kwargs):

        auth = request.headers.get("Authorization", None)
        if auth:
            try:
                _, s = auth.split(" ")
                client_id, client_secret = decode_base64(s).split(":")
                client_id = to_unicode(client_id, "utf-8")
            except Exception as e:
                log.debug("Authenticate client failed with exception: %r", e)
                return False
        else:
            client_id = request.client_id

        client = self._clientgetter(client_id)
        if not client:
            log.debug("Authenticate client failed, client not found.")
            return False

        return self.authenticate_client_id(client_id, request)


oauth = CustomProvider()
oauth._validator = CustomRequestValidator()


@oauth.blueprint.route("/oauth/token", methods=["POST"])
@oauth.token_handler
def access_token(*args, **kwargs):
    return None


@oauth.blueprint.route("/oauth/revoke", methods=["POST"])
@oauth.revoke_handler
def revoke_token():
    pass


@oauth.invalid_response
def invalid_require_oauth(req):
    message = req.error_message if req else "Unauthorized"
    return jsonify(error="invalid_token", message=message), 401

{'fraction_non_alphanumeric': 0.09022681787858572, 'fraction_numerical': 0.004169446297531688, 'mean_word_length': 2.470196759259259, 'pattern_counts': {'":': 10, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5248022', 'n_tokens_mistral': 3371, 'n_tokens_neox': 3205, 'n_words': 795}