Underground-Digital/Workflow-Engine

1import urllib.parse2from dataclasses import dataclass3from typing import Optional4 5import requests6 7 8@dataclass9class OAuthUserInfo:10    id: str11    name: str12    email: str13 14 15class OAuth:16    def __init__(self, client_id: str, client_secret: str, redirect_uri: str):17        self.client_id = client_id18        self.client_secret = client_secret19        self.redirect_uri = redirect_uri20 21    def get_authorization_url(self):22        raise NotImplementedError()23 24    def get_access_token(self, code: str):25        raise NotImplementedError()26 27    def get_raw_user_info(self, token: str):28        raise NotImplementedError()29 30    def get_user_info(self, token: str) -> OAuthUserInfo:31        raw_info = self.get_raw_user_info(token)32        return self._transform_user_info(raw_info)33 34    def _transform_user_info(self, raw_info: dict) -> OAuthUserInfo:35        raise NotImplementedError()36 37 38class GitHubOAuth(OAuth):39    _AUTH_URL = "https://github.com/login/oauth/authorize"40    _TOKEN_URL = "https://github.com/login/oauth/access_token"41    _USER_INFO_URL = "https://api.github.com/user"42    _EMAIL_INFO_URL = "https://api.github.com/user/emails"43 44    def get_authorization_url(self, invite_token: Optional[str] = None):45        params = {46            "client_id": self.client_id,47            "redirect_uri": self.redirect_uri,48            "scope": "user:email",  # Request only basic user information49        }50        if invite_token:51            params["state"] = invite_token52        return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}"53 54    def get_access_token(self, code: str):55        data = {56            "client_id": self.client_id,57            "client_secret": self.client_secret,58            "code": code,59            "redirect_uri": self.redirect_uri,60        }61        headers = {"Accept": "application/json"}62        response = requests.post(self._TOKEN_URL, data=data, headers=headers)63 64        response_json = response.json()65        access_token = response_json.get("access_token")66 67        if not access_token:68            raise ValueError(f"Error in GitHub OAuth: {response_json}")69 70        return access_token71 72    def get_raw_user_info(self, token: str):73        headers = {"Authorization": f"token {token}"}74        response = requests.get(self._USER_INFO_URL, headers=headers)75        response.raise_for_status()76        user_info = response.json()77 78        email_response = requests.get(self._EMAIL_INFO_URL, headers=headers)79        email_info = email_response.json()80        primary_email = next((email for email in email_info if email["primary"] == True), None)81 82        return {**user_info, "email": primary_email["email"]}83 84    def _transform_user_info(self, raw_info: dict) -> OAuthUserInfo:85        email = raw_info.get("email")86        if not email:87            email = f"{raw_info['id']}+{raw_info['login']}@users.noreply.github.com"88        return OAuthUserInfo(id=str(raw_info["id"]), name=raw_info["name"], email=email)89 90 91class GoogleOAuth(OAuth):92    _AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"93    _TOKEN_URL = "https://oauth2.googleapis.com/token"94    _USER_INFO_URL = "https://www.googleapis.com/oauth2/v3/userinfo"95 96    def get_authorization_url(self, invite_token: Optional[str] = None):97        params = {98            "client_id": self.client_id,99            "response_type": "code",100            "redirect_uri": self.redirect_uri,101            "scope": "openid email",102        }103        if invite_token:104            params["state"] = invite_token105        return f"{self._AUTH_URL}?{urllib.parse.urlencode(params)}"106 107    def get_access_token(self, code: str):108        data = {109            "client_id": self.client_id,110            "client_secret": self.client_secret,111            "code": code,112            "grant_type": "authorization_code",113            "redirect_uri": self.redirect_uri,114        }115        headers = {"Accept": "application/json"}116        response = requests.post(self._TOKEN_URL, data=data, headers=headers)117 118        response_json = response.json()119        access_token = response_json.get("access_token")120 121        if not access_token:122            raise ValueError(f"Error in Google OAuth: {response_json}")123 124        return access_token125 126    def get_raw_user_info(self, token: str):127        headers = {"Authorization": f"Bearer {token}"}128        response = requests.get(self._USER_INFO_URL, headers=headers)129        response.raise_for_status()130        return response.json()131 132    def _transform_user_info(self, raw_info: dict) -> OAuthUserInfo:133        return OAuthUserInfo(id=str(raw_info["sub"]), name=None, email=raw_info["email"])134