Skip to content

Commit 78e712d

Browse files
yaozheng-fangyaozheng-fang
committed
feat(auth): add auto auth module
1 parent ecc27d7 commit 78e712d

16 files changed

+1026
-1
lines changed

‎veadk/auth/__init__.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.

‎veadk/auth/base_auth.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
16+
class BaseAuth:
17+
def __init__(self) -> None: ...
18+
19+
def _fetch_token(self) -> str | dict: ...
20+
21+
@property
22+
def token(self) -> str | dict: ...

‎veadk/auth/veauth/__init__.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.

‎veadk/auth/veauth/apmplus_veauth.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
17+
from typing_extensions import override
18+
19+
from veadk.auth.veauth.base_veauth import BaseVeAuth
20+
from veadk.utils.logger import get_logger
21+
from veadk.utils.volcengine_sign import ve_request
22+
23+
logger = get_logger(__name__)
24+
25+
26+
class APMPlusVeAuth(BaseVeAuth):
27+
def __init__(
28+
self,
29+
access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""),
30+
secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""),
31+
) -> None:
32+
super().__init__(access_key, secret_key)
33+
34+
self._token: str = ""
35+
36+
@override
37+
def _fetch_token(self) -> None:
38+
logger.info("Fetching APMPlus token...")
39+
40+
res = ve_request(
41+
request_body={},
42+
action="GetAppKey",
43+
ak=self.access_key,
44+
sk=self.secret_key,
45+
service="apmplus_server",
46+
version="2024-07-30",
47+
region="cn-beijing",
48+
host="open.volcengineapi.com",
49+
)
50+
try:
51+
self._token = res["data"]["app_key"]
52+
except KeyError:
53+
raise ValueError(f"Failed to get APMPlus token: {res}")
54+
55+
@property
56+
def token(self) -> str:
57+
if self._token:
58+
return self._token
59+
self._fetch_token()
60+
return self._token

‎veadk/auth/veauth/ark_veauth.py‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,76 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
17+
from typing_extensions import override
18+
19+
from veadk.auth.veauth.base_veauth import BaseVeAuth
20+
from veadk.utils.logger import get_logger
21+
from veadk.utils.volcengine_sign import ve_request
22+
23+
logger = get_logger(__name__)
24+
25+
26+
class ARKVeAuth(BaseVeAuth):
27+
def __init__(
28+
self,
29+
access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""),
30+
secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""),
31+
) -> None:
32+
super().__init__(access_key, secret_key)
33+
34+
self._token: str = ""
35+
36+
@override
37+
def _fetch_token(self) -> None:
38+
# list api keys
39+
first_api_key_id = ""
40+
res = ve_request(
41+
request_body={"ProjectName": "default", "Filter": {}},
42+
action="ListApiKeys",
43+
ak=self.access_key,
44+
sk=self.secret_key,
45+
service="ark",
46+
version="2024-01-01",
47+
region="cn-beijing",
48+
host="open.volcengineapi.com",
49+
)
50+
try:
51+
first_api_key_id = res["Result"]["Items"][0]["Id"]
52+
except KeyError:
53+
raise ValueError(f"Failed to get ARK api key list: {res}")
54+
55+
# get raw api key
56+
res = ve_request(
57+
request_body={"Id": first_api_key_id},
58+
action="GetRawApiKey",
59+
ak=self.access_key,
60+
sk=self.secret_key,
61+
service="ark",
62+
version="2024-01-01",
63+
region="cn-beijing",
64+
host="open.volcengineapi.com",
65+
)
66+
try:
67+
self._token = res["Result"]["ApiKey"]
68+
except KeyError:
69+
raise ValueError(f"Failed to get ARK api key: {res}")
70+
71+
@property
72+
def token(self) -> str:
73+
if self._token:
74+
return self._token
75+
self._fetch_token()
76+
return self._token

‎veadk/auth/veauth/base_veauth.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
from abc import ABC, abstractmethod
17+
from typing import Type
18+
19+
from veadk.auth.base_auth import BaseAuth
20+
21+
22+
class BaseVeAuth(ABC, BaseAuth):
23+
volcengine_access_key: str
24+
"""Volcengine Access Key"""
25+
26+
volcengine_secret_key: str
27+
"""Volcengine Secret Key"""
28+
29+
def __init__(
30+
self,
31+
access_key: str | None = None,
32+
secret_key: str | None = None,
33+
) -> None:
34+
super().__init__()
35+
36+
final_ak = access_key or os.getenv("VOLCENGINE_ACCESS_KEY")
37+
final_sk = secret_key or os.getenv("VOLCENGINE_SECRET_KEY")
38+
39+
assert final_ak, "Volcengine access key cannot be empty."
40+
assert final_sk, "Volcengine secret key cannot be empty."
41+
42+
self.access_key = final_ak
43+
self.secret_key = final_sk
44+
45+
self._token: str = ""
46+
47+
@abstractmethod
48+
def _fetch_token(self) -> None: ...
49+
50+
@property
51+
def token(self) -> str: ...
52+
53+
54+
def veauth(auth_token_name: str, auth_cls: Type[BaseVeAuth]):
55+
def decorator(cls: Type):
56+
# api_key -> _api_key
57+
# for cache
58+
private_auth_token_name = f"_{auth_token_name}"
59+
setattr(cls, private_auth_token_name, "")
60+
61+
# init a auth cls for fetching token
62+
auth_cls_instance = "_auth_cls_instance"
63+
setattr(cls, auth_cls_instance, auth_cls())
64+
65+
def getattribute(self, name: str):
66+
if name != auth_token_name:
67+
return object.__getattribute__(self, name)
68+
if name == auth_token_name:
69+
token = object.__getattribute__(self, name)
70+
71+
if token:
72+
return token
73+
elif not token and not getattr(cls, private_auth_token_name):
74+
token = getattr(cls, auth_cls_instance).token
75+
setattr(cls, private_auth_token_name, token)
76+
return token
77+
elif not token and getattr(cls, private_auth_token_name):
78+
return getattr(cls, private_auth_token_name)
79+
return token
80+
81+
setattr(cls, "__getattribute__", getattribute)
82+
return cls
83+
84+
return decorator

‎veadk/auth/veauth/cozeloop_veauth.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.

‎veadk/auth/veauth/prompt_pilot_veauth.py‎

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
import os
16+
17+
from typing_extensions import override
18+
19+
from veadk.auth.veauth.base_veauth import BaseVeAuth
20+
from veadk.utils.logger import get_logger
21+
from veadk.utils.volcengine_sign import ve_request
22+
23+
logger = get_logger(__name__)
24+
25+
26+
class PromptPilotVeAuth(BaseVeAuth):
27+
def __init__(
28+
self,
29+
access_key: str = os.getenv("VOLCENGINE_ACCESS_KEY", ""),
30+
secret_key: str = os.getenv("VOLCENGINE_SECRET_KEY", ""),
31+
) -> None:
32+
super().__init__(access_key, secret_key)
33+
34+
self._token: str | dict = ""
35+
36+
@override
37+
def _fetch_token(self) -> None:
38+
logger.info("Fetching Prompt Pilot token...")
39+
40+
res = ve_request(
41+
request_body={},
42+
action="GetOrCreatePromptPilotAPIKeys",
43+
ak=self.access_key,
44+
sk=self.secret_key,
45+
service="ark",
46+
version="2024-01-01",
47+
region="cn-beijing",
48+
host="open.volcengineapi.com",
49+
)
50+
try:
51+
self._token = res["Result"]["APIKeys"][0]["APIKey"]
52+
except KeyError:
53+
raise ValueError(f"Failed to get Prompt Pilot token: {res}")
54+
55+
@property
56+
def token(self) -> str | dict:
57+
if self._token:
58+
return self._token
59+
self._fetch_token()
60+
return self._token

0 commit comments

Comments
 (0)