mirror of https://github.com/jumpserver/jumpserver
150 lines
4.5 KiB
Python
150 lines
4.5 KiB
Python
import requests
|
||
from requests import exceptions as req_exce
|
||
from rest_framework.exceptions import PermissionDenied
|
||
from django.core.cache import cache
|
||
|
||
from .utils import DictWrapper
|
||
from common.utils.common import get_logger
|
||
from common.utils import lazyproperty
|
||
from common.message.backends.utils import set_default, as_request
|
||
|
||
from . import exceptions as exce
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
|
||
class RequestMixin:
|
||
code_key: str
|
||
msg_key: str
|
||
|
||
|
||
class BaseRequest(RequestMixin):
|
||
"""
|
||
定义了 `access_token` 的过期刷新框架
|
||
"""
|
||
invalid_token_errcodes = ()
|
||
code_key = 'errcode'
|
||
msg_key = 'err_msg'
|
||
|
||
def __init__(self, timeout=None):
|
||
self._request_kwargs = {
|
||
'timeout': timeout
|
||
}
|
||
self.init_access_token()
|
||
|
||
@classmethod
|
||
def check_errcode_is_0(cls, data: DictWrapper):
|
||
errcode = data[cls.code_key]
|
||
if errcode != 0:
|
||
# 如果代码写的对，配置没问题，这里不该出错，系统性错误，直接抛异常
|
||
errmsg = data[cls.msg_key]
|
||
logger.error(f'Response 200 but errcode is not 0: '
|
||
f'errcode={errcode} '
|
||
f'errmsg={errmsg} ')
|
||
raise exce.ErrCodeNot0(detail=data.raw_data)
|
||
|
||
@staticmethod
|
||
def check_http_is_200(response):
|
||
if response.status_code != 200:
|
||
# 正常情况下不会返回非 200 响应码
|
||
logger.error(f'Response error: '
|
||
f'status_code={response.status_code} '
|
||
f'url={response.url}'
|
||
f'\ncontent={response.content}')
|
||
raise exce.HTTPNot200(detail=response.json())
|
||
|
||
def request_access_token(self):
|
||
"""
|
||
获取新的 `access_token` 的方法，子类需要实现
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
def get_access_token_cache_key(self):
|
||
"""
|
||
获取 `access_token` 的缓存 key， 子类需要实现
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
def add_token(self, kwargs: dict):
|
||
"""
|
||
添加 token ，子类需要实现
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
def is_token_invalid(self, data):
|
||
code = data[self.code_key]
|
||
if code in self.invalid_token_errcodes:
|
||
logger.error(f'OAuth token invalid: {data}')
|
||
return True
|
||
return False
|
||
|
||
@lazyproperty
|
||
def access_token_cache_key(self):
|
||
return self.get_access_token_cache_key()
|
||
|
||
def init_access_token(self):
|
||
access_token = cache.get(self.access_token_cache_key)
|
||
if access_token:
|
||
self.access_token = access_token
|
||
return
|
||
self.refresh_access_token()
|
||
|
||
def refresh_access_token(self):
|
||
access_token, expires_in = self.request_access_token()
|
||
self.access_token = access_token
|
||
cache.set(self.access_token_cache_key, access_token, expires_in - 10)
|
||
|
||
def raw_request(self, method, url, **kwargs):
|
||
set_default(kwargs, self._request_kwargs)
|
||
try:
|
||
response = getattr(requests, method)(url, **kwargs)
|
||
self.check_http_is_200(response)
|
||
raw_data = response.json()
|
||
data = DictWrapper(raw_data)
|
||
|
||
return data
|
||
except req_exce.ReadTimeout as e:
|
||
logger.exception(e)
|
||
raise exce.NetError
|
||
|
||
def token_request(self, method, url, **kwargs):
|
||
for i in range(3):
|
||
# 循环为了防止 access_token 失效
|
||
self.add_token(kwargs)
|
||
data = self.raw_request(method, url, **kwargs)
|
||
|
||
if self.is_token_invalid(data):
|
||
self.refresh_access_token()
|
||
continue
|
||
|
||
return data
|
||
logger.error(f'Get access_token error, check config: url={url} data={data.raw_data}')
|
||
raise PermissionDenied(data.raw_data)
|
||
|
||
def get(self, url, params=None, with_token=True,
|
||
check_errcode_is_0=True, **kwargs):
|
||
# self.request ...
|
||
pass
|
||
get = as_request(get)
|
||
|
||
def post(self, url, params=None, json=None,
|
||
with_token=True, check_errcode_is_0=True,
|
||
**kwargs):
|
||
# self.request ...
|
||
pass
|
||
post = as_request(post)
|
||
|
||
def request(self, method, url,
|
||
with_token=True,
|
||
check_errcode_is_0=True,
|
||
**kwargs):
|
||
|
||
if with_token:
|
||
data = self.token_request(method, url, **kwargs)
|
||
else:
|
||
data = self.raw_request(method, url, **kwargs)
|
||
|
||
if check_errcode_is_0:
|
||
self.check_errcode_is_0(data)
|
||
return data
|