#!/usr/bin/env python
# coding: utf-8
"""
OAuth2 提供商
提供主流 OAuth2 提供商的实现，包括 GitHub、Google、微信等。
"""
import base64
import hashlib
import secrets
import urllib.parse
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any, Dict, Optional
from .._json import loads
[文档]
@dataclass
class OAuth2UserInfo:
"""OAuth2 用户信息"""
provider: str
provider_user_id: str
username: str
email: Optional[str] = None
name: Optional[str] = None
avatar_url: Optional[str] = None
raw_data: Optional[Dict[str, Any]] = None
[文档]
def to_dict(self) -> Dict[str, Any]:
return {
'provider': self.provider,
'provider_user_id': self.provider_user_id,
'username': self.username,
'email': self.email,
'name': self.name,
'avatar_url': self.avatar_url,
}
[文档]
class OAuth2Provider(ABC):
"""OAuth2 提供商基类"""
name: str = 'base'
authorize_url: str = ''
access_token_url: str = ''
userinfo_url: str = ''
scope: str = ''
def __init__(
self,
client_id: str,
client_secret: str,
redirect_uri: str = '',
scope: str = None,
**kwargs
):
self.client_id = client_id
self.client_secret = client_secret
self.redirect_uri = redirect_uri
self.scope = scope or self.scope
self.extra_params = kwargs
[文档]
def get_authorize_url(
self,
state: str,
redirect_uri: str = None,
**kwargs
) -> str:
"""
生成授权 URL
Args:
state: 状态参数，用于防止 CSRF 攻击
redirect_uri: 回调地址
**kwargs: 额外参数
Returns:
授权 URL
"""
params = {
'client_id': self.client_id,
'redirect_uri': redirect_uri or self.redirect_uri,
'response_type': 'code',
'state': state,
}
if self.scope:
params['scope'] = self.scope
params.update(kwargs)
return f"{self.authorize_url}?{urllib.parse.urlencode(params)}"
[文档]
@abstractmethod
def get_access_token(
self,
code: str,
redirect_uri: str = None
) -> Dict[str, Any]:
"""
使用授权码获取 Access Token
Args:
code: 授权码
redirect_uri: 回调地址
Returns:
Token 信息
"""
pass
[文档]
@abstractmethod
def get_user_info(self, access_token: str) -> OAuth2UserInfo:
"""
获取用户信息
Args:
access_token: Access Token
Returns:
用户信息
"""
pass
def _request(
self,
url: str,
method: str = 'GET',
data: Dict = None,
headers: Dict = None
) -> Dict[str, Any]:
"""
发送 HTTP 请求（同步）
Args:
url: 请求 URL
method: 请求方法
data: 请求数据
headers: 请求头
Returns:
响应数据
"""
try:
import urllib.request
import urllib.error
req_headers = {
'Accept': 'application/json',
'User-Agent': 'Litefs-OAuth2/1.0',
}
if headers:
req_headers.update(headers)
if method == 'GET':
if data:
url = f"{url}?{urllib.parse.urlencode(data)}"
req = urllib.request.Request(url, headers=req_headers)
else:
body = urllib.parse.urlencode(data).encode('utf-8') if data else b''
req = urllib.request.Request(url, data=body, headers=req_headers, method=method)
with urllib.request.urlopen(req, timeout=30) as resp:
return loads(resp.read())
except urllib.error.HTTPError as e:
error_body = e.read().decode('utf-8') if e.fp else ''
return {'error': f'HTTP {e.code}', 'error_description': error_body}
except Exception as e:
return {'error': str(e)}
async def _async_request(
self,
url: str,
method: str = 'GET',
data: Dict = None,
headers: Dict = None
) -> Dict[str, Any]:
"""
发送 HTTP 请求（异步）
需要 aiohttp 库，如果未安装则回退到同步请求。
Args:
url: 请求 URL
method: 请求方法
data: 请求数据
headers: 请求头
Returns:
响应数据
"""
import logging
try:
import aiohttp
except ImportError:
logging.getLogger(__name__).warning(
"aiohttp is not installed, falling back to synchronous OAuth2 request. "
"Install aiohttp for async support: pip install aiohttp"
)
return self._request(url, method=method, data=data, headers=headers)
req_headers = {
'Accept': 'application/json',
'User-Agent': 'Litefs-OAuth2/1.0',
}
if headers:
req_headers.update(headers)
try:
async with aiohttp.ClientSession() as session:
if method == 'GET':
if data:
url = f"{url}?{urllib.parse.urlencode(data)}"
async with session.get(
url, headers=req_headers, timeout=aiohttp.ClientTimeout(total=30)
) as resp:
return await resp.json(content_type=None)
else:
async with session.request(
method, url, data=data, headers=req_headers,
timeout=aiohttp.ClientTimeout(total=30)
) as resp:
return await resp.json(content_type=None)
except Exception as e:
return {'error': str(e)}
[文档]
def generate_state(self) -> str:
"""生成随机状态参数"""
return secrets.token_urlsafe(32)
[文档]
def generate_pkce_verifier(self) -> str:
"""生成 PKCE code_verifier"""
return secrets.token_urlsafe(64)
[文档]
def generate_pkce_challenge(self, verifier: str, method: str = 'S256') -> str:
"""
生成 PKCE code_challenge
Args:
verifier: code_verifier
method: 挑战方法 ('S256' 或 'plain')
Returns:
code_challenge
"""
if method == 'S256':
digest = hashlib.sha256(verifier.encode('utf-8')).digest()
return base64.urlsafe_b64encode(digest).decode('utf-8').rstrip('=')
return verifier
[文档]
class GitHubProvider(OAuth2Provider):
"""GitHub OAuth2 提供商"""
name = 'github'
authorize_url = 'https://github.com/login/oauth/authorize'
access_token_url = 'https://github.com/login/oauth/access_token'
userinfo_url = 'https://api.github.com/user'
scope = 'user:email'
[文档]
def get_access_token(
self,
code: str,
redirect_uri: str = None
) -> Dict[str, Any]:
data = {
'client_id': self.client_id,
'client_secret': self.client_secret,
'code': code,
'redirect_uri': redirect_uri or self.redirect_uri,
}
result = self._request(
self.access_token_url,
method='POST',
data=data,
headers={'Accept': 'application/json'}
)
if 'error' in result:
return result
return result
[文档]
def get_user_info(self, access_token: str) -> OAuth2UserInfo:
headers = {
'Authorization': f'Bearer {access_token}',
'Accept': 'application/vnd.github.v3+json',
}
user_data = self._request(self.userinfo_url, headers=headers)
if 'error' in user_data:
return None
emails = self._request(
'https://api.github.com/user/emails',
headers=headers
)
primary_email = None
if isinstance(emails, list):
for email_info in emails:
if email_info.get('primary') and email_info.get('verified'):
primary_email = email_info.get('email')
break
return OAuth2UserInfo(
provider=self.name,
provider_user_id=str(user_data.get('id', '')),
username=user_data.get('login', ''),
email=primary_email or user_data.get('email'),
name=user_data.get('name') or user_data.get('login'),
avatar_url=user_data.get('avatar_url'),
raw_data=user_data,
)
[文档]
class GoogleProvider(OAuth2Provider):
"""Google OAuth2 提供商"""
name = 'google'
authorize_url = 'https://accounts.google.com/o/oauth2/v2/auth'
access_token_url = 'https://oauth2.googleapis.com/token'
userinfo_url = 'https://www.googleapis.com/oauth2/v3/userinfo'
scope = 'openid email profile'
[文档]
def get_access_token(
self,
code: str,
redirect_uri: str = None
) -> Dict[str, Any]:
data = {
'client_id': self.client_id,
'client_secret': self.client_secret,
'code': code,
'redirect_uri': redirect_uri or self.redirect_uri,
'grant_type': 'authorization_code',
}
result = self._request(
self.access_token_url,
method='POST',
data=data
)
return result
[文档]
def get_user_info(self, access_token: str) -> OAuth2UserInfo:
headers = {
'Authorization': f'Bearer {access_token}',
}
user_data = self._request(self.userinfo_url, headers=headers)
if 'error' in user_data:
return None
return OAuth2UserInfo(
provider=self.name,
provider_user_id=str(user_data.get('sub', '')),
username=user_data.get('email', '').split('@')[0],
email=user_data.get('email'),
name=user_data.get('name'),
avatar_url=user_data.get('picture'),
raw_data=user_data,
)
[文档]
class WeChatProvider(OAuth2Provider):
"""微信 OAuth2 提供商"""
name = 'wechat'
authorize_url = 'https://open.weixin.qq.com/connect/qrconnect'
access_token_url = 'https://api.weixin.qq.com/sns/oauth2/access_token'
userinfo_url = 'https://api.weixin.qq.com/sns/userinfo'
scope = 'snsapi_login'
[文档]
def get_authorize_url(
self,
state: str,
redirect_uri: str = None,
**kwargs
) -> str:
params = {
'appid': self.client_id,
'redirect_uri': redirect_uri or self.redirect_uri,
'response_type': 'code',
'scope': self.scope,
'state': state,
}
return f"{self.authorize_url}?{urllib.parse.urlencode(params)}#wechat_redirect"
[文档]
def get_access_token(
self,
code: str,
redirect_uri: str = None
) -> Dict[str, Any]:
params = {
'appid': self.client_id,
'secret': self.client_secret,
'code': code,
'grant_type': 'authorization_code',
}
url = f"{self.access_token_url}?{urllib.parse.urlencode(params)}"
result = self._request(url)
return result
[文档]
def get_user_info(self, access_token: str, openid: str = None) -> OAuth2UserInfo:
if not openid:
return None
params = {
'access_token': access_token,
'openid': openid,
}
url = f"{self.userinfo_url}?{urllib.parse.urlencode(params)}"
user_data = self._request(url)
if 'errcode' in user_data:
return None
return OAuth2UserInfo(
provider=self.name,
provider_user_id=str(user_data.get('unionid') or user_data.get('openid', '')),
username=user_data.get('nickname', ''),
name=user_data.get('nickname'),
avatar_url=user_data.get('headimgurl'),
raw_data=user_data,
)
[文档]
class WeComProvider(OAuth2Provider):
"""企业微信 OAuth2 提供商"""
name = 'wecom'
authorize_url = 'https://open.work.weixin.qq.com/wwopen/sso/qrConnect'
access_token_url = 'https://qyapi.weixin.qq.com/cgi-bin/gettoken'
userinfo_url = 'https://qyapi.weixin.qq.com/cgi-bin/user/getuserinfo'
scope = 'snsapi_base'
def __init__(
self,
client_id: str,
client_secret: str,
agent_id: str = '',
redirect_uri: str = '',
**kwargs
):
super().__init__(client_id, client_secret, redirect_uri, **kwargs)
self.agent_id = agent_id
[文档]
def get_authorize_url(
self,
state: str,
redirect_uri: str = None,
**kwargs
) -> str:
params = {
'appid': self.client_id,
'agentid': self.agent_id,
'redirect_uri': redirect_uri or self.redirect_uri,
'response_type': 'code',
'scope': self.scope,
'state': state,
}
return f"{self.authorize_url}?{urllib.parse.urlencode(params)}"
[文档]
def get_access_token(
self,
code: str,
redirect_uri: str = None
) -> Dict[str, Any]:
params = {
'corpid': self.client_id,
'corpsecret': self.client_secret,
}
url = f"{self.access_token_url}?{urllib.parse.urlencode(params)}"
result = self._request(url)
if 'errcode' in result and result['errcode'] != 0:
return result
return result
[文档]
def get_user_info(self, access_token: str, code: str = None) -> OAuth2UserInfo:
if not code:
return None
params = {
'access_token': access_token,
'code': code,
}
url = f"{self.userinfo_url}?{urllib.parse.urlencode(params)}"
user_data = self._request(url)
if 'errcode' in user_data and user_data['errcode'] != 0:
return None
return OAuth2UserInfo(
provider=self.name,
provider_user_id=str(user_data.get('UserId', '')),
username=user_data.get('UserId', ''),
name=user_data.get('name'),
raw_data=user_data,
)
PROVIDERS = {
'github': GitHubProvider,
'google': GoogleProvider,
'wechat': WeChatProvider,
'wecom': WeComProvider,
}
[文档]
def get_provider(name: str) -> Optional[type]:
"""获取提供商类"""
return PROVIDERS.get(name.lower())