ya-oauth 0.74.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ya_oauth-0.74.0/.gitignore +169 -0
- ya_oauth-0.74.0/PKG-INFO +32 -0
- ya_oauth-0.74.0/README.md +13 -0
- ya_oauth-0.74.0/pyproject.toml +58 -0
- ya_oauth-0.74.0/tests/test_codex.py +155 -0
- ya_oauth-0.74.0/tests/test_store.py +18 -0
- ya_oauth-0.74.0/ya_oauth/__init__.py +15 -0
- ya_oauth-0.74.0/ya_oauth/cli.py +146 -0
- ya_oauth-0.74.0/ya_oauth/codex.py +239 -0
- ya_oauth-0.74.0/ya_oauth/jwt.py +50 -0
- ya_oauth-0.74.0/ya_oauth/store.py +161 -0
- ya_oauth-0.74.0/ya_oauth/types.py +87 -0
|
@@ -0,0 +1,169 @@
|
|
|
1
|
+
docs/source
|
|
2
|
+
|
|
3
|
+
# From https://raw.githubusercontent.com/github/gitignore/main/Python.gitignore
|
|
4
|
+
|
|
5
|
+
# Byte-compiled / optimized / DLL files
|
|
6
|
+
__pycache__/
|
|
7
|
+
*.py[cod]
|
|
8
|
+
*$py.class
|
|
9
|
+
|
|
10
|
+
# C extensions
|
|
11
|
+
*.so
|
|
12
|
+
|
|
13
|
+
# Distribution / packaging
|
|
14
|
+
.Python
|
|
15
|
+
build/
|
|
16
|
+
develop-eggs/
|
|
17
|
+
dist/
|
|
18
|
+
downloads/
|
|
19
|
+
eggs/
|
|
20
|
+
.eggs/
|
|
21
|
+
lib/
|
|
22
|
+
!apps/
|
|
23
|
+
!apps/ya-claw-web/
|
|
24
|
+
!apps/ya-claw-web/src/
|
|
25
|
+
!apps/ya-claw-web/src/lib/
|
|
26
|
+
!apps/ya-claw-web/src/lib/**
|
|
27
|
+
lib64/
|
|
28
|
+
parts/
|
|
29
|
+
sdist/
|
|
30
|
+
var/
|
|
31
|
+
wheels/
|
|
32
|
+
share/python-wheels/
|
|
33
|
+
*.egg-info/
|
|
34
|
+
.installed.cfg
|
|
35
|
+
*.egg
|
|
36
|
+
MANIFEST
|
|
37
|
+
|
|
38
|
+
# PyInstaller
|
|
39
|
+
# Usually these files are written by a python script from a template
|
|
40
|
+
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
|
41
|
+
*.manifest
|
|
42
|
+
*.spec
|
|
43
|
+
|
|
44
|
+
# Installer logs
|
|
45
|
+
pip-log.txt
|
|
46
|
+
pip-delete-this-directory.txt
|
|
47
|
+
|
|
48
|
+
# Unit test / coverage reports
|
|
49
|
+
htmlcov/
|
|
50
|
+
.tox/
|
|
51
|
+
.nox/
|
|
52
|
+
.coverage
|
|
53
|
+
.coverage.*
|
|
54
|
+
.cache
|
|
55
|
+
nosetests.xml
|
|
56
|
+
coverage.xml
|
|
57
|
+
*.cover
|
|
58
|
+
*.py,cover
|
|
59
|
+
.hypothesis/
|
|
60
|
+
.pytest_cache/
|
|
61
|
+
cover/
|
|
62
|
+
|
|
63
|
+
# Translations
|
|
64
|
+
*.mo
|
|
65
|
+
*.pot
|
|
66
|
+
|
|
67
|
+
# Django stuff:
|
|
68
|
+
*.log
|
|
69
|
+
local_settings.py
|
|
70
|
+
db.sqlite3
|
|
71
|
+
db.sqlite3-journal
|
|
72
|
+
|
|
73
|
+
# Flask stuff:
|
|
74
|
+
instance/
|
|
75
|
+
.webassets-cache
|
|
76
|
+
|
|
77
|
+
# Scrapy stuff:
|
|
78
|
+
.scrapy
|
|
79
|
+
|
|
80
|
+
# Sphinx documentation
|
|
81
|
+
docs/_build/
|
|
82
|
+
|
|
83
|
+
# PyBuilder
|
|
84
|
+
.pybuilder/
|
|
85
|
+
target/
|
|
86
|
+
|
|
87
|
+
# Jupyter Notebook
|
|
88
|
+
.ipynb_checkpoints
|
|
89
|
+
|
|
90
|
+
# IPython
|
|
91
|
+
profile_default/
|
|
92
|
+
ipython_config.py
|
|
93
|
+
|
|
94
|
+
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
|
95
|
+
__pypackages__/
|
|
96
|
+
|
|
97
|
+
# Celery stuff
|
|
98
|
+
celerybeat-schedule
|
|
99
|
+
celerybeat.pid
|
|
100
|
+
|
|
101
|
+
# SageMath parsed files
|
|
102
|
+
*.sage.py
|
|
103
|
+
|
|
104
|
+
# Environments
|
|
105
|
+
.env
|
|
106
|
+
.yaai/
|
|
107
|
+
**/.yaai/
|
|
108
|
+
.venv
|
|
109
|
+
env/
|
|
110
|
+
venv/
|
|
111
|
+
ENV/
|
|
112
|
+
env.bak/
|
|
113
|
+
venv.bak/
|
|
114
|
+
|
|
115
|
+
# Spyder project settings
|
|
116
|
+
.spyderproject
|
|
117
|
+
.spyproject
|
|
118
|
+
|
|
119
|
+
# Rope project settings
|
|
120
|
+
.ropeproject
|
|
121
|
+
|
|
122
|
+
# mkdocs documentation
|
|
123
|
+
/site
|
|
124
|
+
|
|
125
|
+
# mypy
|
|
126
|
+
.mypy_cache/
|
|
127
|
+
.pyright/
|
|
128
|
+
.dmypy.json
|
|
129
|
+
dmypy.json
|
|
130
|
+
|
|
131
|
+
# Pyre type checker
|
|
132
|
+
.pyre/
|
|
133
|
+
|
|
134
|
+
# pytype static type analyzer
|
|
135
|
+
.pytype/
|
|
136
|
+
|
|
137
|
+
# Cython debug symbols
|
|
138
|
+
cython_debug/
|
|
139
|
+
|
|
140
|
+
# Vscode config files
|
|
141
|
+
# .vscode/
|
|
142
|
+
|
|
143
|
+
# PyCharm
|
|
144
|
+
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
|
145
|
+
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
|
146
|
+
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
|
147
|
+
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
|
148
|
+
#.idea/
|
|
149
|
+
|
|
150
|
+
TO-DO.json
|
|
151
|
+
dev/
|
|
152
|
+
ya_agent_sdk/sandbox/shell/templates/public
|
|
153
|
+
|
|
154
|
+
# Sync automaticlly
|
|
155
|
+
yaacli/yaacli/skills/building-agents
|
|
156
|
+
|
|
157
|
+
# Frontend
|
|
158
|
+
node_modules/
|
|
159
|
+
apps/*/dist/
|
|
160
|
+
!apps/ya-desktop/
|
|
161
|
+
!apps/ya-desktop/src/
|
|
162
|
+
!apps/ya-desktop/src/**
|
|
163
|
+
!apps/ya-desktop/src-tauri/
|
|
164
|
+
!apps/ya-desktop/src-tauri/resources/
|
|
165
|
+
!apps/ya-desktop/src-tauri/resources/uv/
|
|
166
|
+
!apps/ya-desktop/src-tauri/resources/uv/.gitkeep
|
|
167
|
+
apps/ya-desktop/src-tauri/resources/uv/uv
|
|
168
|
+
apps/ya-desktop/src-tauri/resources/uv/uv.exe
|
|
169
|
+
*.tsbuildinfo
|
ya_oauth-0.74.0/PKG-INFO
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: ya-oauth
|
|
3
|
+
Version: 0.74.0
|
|
4
|
+
Summary: OAuth login, refresh, storage, and CLI for YA model providers
|
|
5
|
+
Project-URL: Repository, https://github.com/wh1isper/ya-mono
|
|
6
|
+
Author-email: wh1isper <jizhongsheng957@gmail.com>
|
|
7
|
+
Keywords: agents,ai,oauth
|
|
8
|
+
Classifier: Intended Audience :: Developers
|
|
9
|
+
Classifier: Programming Language :: Python
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
14
|
+
Requires-Python: <3.14,>=3.11
|
|
15
|
+
Requires-Dist: click>=8.0
|
|
16
|
+
Requires-Dist: httpx>=0.28.1
|
|
17
|
+
Requires-Dist: pydantic>=2.12.0
|
|
18
|
+
Description-Content-Type: text/markdown
|
|
19
|
+
|
|
20
|
+
# ya-oauth
|
|
21
|
+
|
|
22
|
+
OAuth login, refresh, logout, token storage, and CLI for YA model providers.
|
|
23
|
+
|
|
24
|
+
## Codex login
|
|
25
|
+
|
|
26
|
+
```bash
|
|
27
|
+
ya-oauth login codex
|
|
28
|
+
ya-oauth status codex
|
|
29
|
+
ya-oauth refresh codex
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
Credentials are stored in `~/.yaai/auth.json` with directory mode `0700` and file mode `0600`.
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
# ya-oauth
|
|
2
|
+
|
|
3
|
+
OAuth login, refresh, logout, token storage, and CLI for YA model providers.
|
|
4
|
+
|
|
5
|
+
## Codex login
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
ya-oauth login codex
|
|
9
|
+
ya-oauth status codex
|
|
10
|
+
ya-oauth refresh codex
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
Credentials are stored in `~/.yaai/auth.json` with directory mode `0700` and file mode `0600`.
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "ya-oauth"
|
|
3
|
+
dynamic = ["version", "dependencies"]
|
|
4
|
+
description = "OAuth login, refresh, storage, and CLI for YA model providers"
|
|
5
|
+
authors = [{ name = "wh1isper", email = "jizhongsheng957@gmail.com" }]
|
|
6
|
+
readme = "README.md"
|
|
7
|
+
keywords = ["oauth", "ai", "agents"]
|
|
8
|
+
requires-python = ">=3.11,<3.14"
|
|
9
|
+
classifiers = [
|
|
10
|
+
"Intended Audience :: Developers",
|
|
11
|
+
"Programming Language :: Python",
|
|
12
|
+
"Programming Language :: Python :: 3",
|
|
13
|
+
"Programming Language :: Python :: 3.11",
|
|
14
|
+
"Programming Language :: Python :: 3.12",
|
|
15
|
+
"Programming Language :: Python :: 3.13",
|
|
16
|
+
]
|
|
17
|
+
|
|
18
|
+
[project.urls]
|
|
19
|
+
Repository = "https://github.com/wh1isper/ya-mono"
|
|
20
|
+
|
|
21
|
+
[project.scripts]
|
|
22
|
+
ya-oauth = "ya_oauth.cli:cli"
|
|
23
|
+
|
|
24
|
+
[dependency-groups]
|
|
25
|
+
dev = [
|
|
26
|
+
"pytest>=7.2.0",
|
|
27
|
+
"pytest-asyncio>=0.25.3",
|
|
28
|
+
"pytest-httpx>=0.35.0",
|
|
29
|
+
"ruff>=0.9.2",
|
|
30
|
+
"pyright>=1.1.0",
|
|
31
|
+
]
|
|
32
|
+
|
|
33
|
+
[build-system]
|
|
34
|
+
requires = ["hatchling", "uv-dynamic-versioning>=0.7.0"]
|
|
35
|
+
build-backend = "hatchling.build"
|
|
36
|
+
|
|
37
|
+
[tool.hatch.version]
|
|
38
|
+
source = "uv-dynamic-versioning"
|
|
39
|
+
|
|
40
|
+
[tool.uv-dynamic-versioning]
|
|
41
|
+
vcs = "git"
|
|
42
|
+
style = "pep440"
|
|
43
|
+
bump = true
|
|
44
|
+
|
|
45
|
+
[tool.hatch.metadata.hooks.uv-dynamic-versioning]
|
|
46
|
+
dependencies = [
|
|
47
|
+
"click>=8.0",
|
|
48
|
+
"httpx>=0.28.1",
|
|
49
|
+
"pydantic>=2.12.0",
|
|
50
|
+
]
|
|
51
|
+
|
|
52
|
+
[tool.hatch.build.targets.wheel]
|
|
53
|
+
packages = ["ya_oauth"]
|
|
54
|
+
|
|
55
|
+
[tool.deptry]
|
|
56
|
+
ignore = ["DEP001", "DEP002"]
|
|
57
|
+
package_module_name_map = { "click" = "click", "httpx" = "httpx", "pydantic" = "pydantic", "pytest" = "pytest", "pytest-asyncio" = "pytest_asyncio", "pytest-httpx" = "pytest_httpx", "ruff" = "ruff", "pyright" = "pyright" }
|
|
58
|
+
per_rule_ignores = { DEP003 = ["click", "httpx", "pydantic", "pytest"], DEP004 = ["pytest"] }
|
|
@@ -0,0 +1,155 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import base64
|
|
4
|
+
import json
|
|
5
|
+
from datetime import UTC, datetime
|
|
6
|
+
|
|
7
|
+
import httpx
|
|
8
|
+
from ya_oauth.codex import CODEX_CLIENT_ID, CODEX_TOKEN_ENDPOINT, CodexOAuthClient
|
|
9
|
+
from ya_oauth.jwt import account_from_id_token
|
|
10
|
+
from ya_oauth.store import OAuthStore
|
|
11
|
+
from ya_oauth.types import OAuthAccount, OAuthProviderRecord, OAuthTokens
|
|
12
|
+
|
|
13
|
+
ACCESS_TOKEN = "fixture-access-token" # noqa: S105
|
|
14
|
+
REFRESH_TOKEN = "fixture-refresh-token" # noqa: S105
|
|
15
|
+
OLD_ACCESS_TOKEN = "fixture-old-access-token" # noqa: S105
|
|
16
|
+
OLD_REFRESH_TOKEN = "fixture-old-refresh-token" # noqa: S105
|
|
17
|
+
NEW_ACCESS_TOKEN = "fixture-new-access-token" # noqa: S105
|
|
18
|
+
OLD_ID_TOKEN = "fixture-old-id-token" # noqa: S105
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def _jwt(payload: dict[str, object]) -> str:
|
|
22
|
+
def enc(data: dict[str, object]) -> str:
|
|
23
|
+
raw = json.dumps(data, separators=(",", ":")).encode()
|
|
24
|
+
return base64.urlsafe_b64encode(raw).decode().rstrip("=")
|
|
25
|
+
|
|
26
|
+
return f"{enc({'alg': 'none'})}.{enc(payload)}.sig"
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def test_account_from_id_token_parses_codex_claims() -> None:
|
|
30
|
+
token = _jwt({
|
|
31
|
+
"email": "top@example.com",
|
|
32
|
+
"https://api.openai.com/auth": {
|
|
33
|
+
"chatgpt_plan_type": "plus",
|
|
34
|
+
"chatgpt_user_id": "user_123",
|
|
35
|
+
"chatgpt_account_id": "acct_123",
|
|
36
|
+
"chatgpt_account_is_fedramp": True,
|
|
37
|
+
},
|
|
38
|
+
})
|
|
39
|
+
|
|
40
|
+
account = account_from_id_token(token)
|
|
41
|
+
|
|
42
|
+
assert account.email == "top@example.com"
|
|
43
|
+
assert account.chatgpt_plan_type == "plus"
|
|
44
|
+
assert account.chatgpt_user_id == "user_123"
|
|
45
|
+
assert account.chatgpt_account_id == "acct_123"
|
|
46
|
+
assert account.chatgpt_account_is_fedramp is True
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
def test_codex_device_code_login_requests_match_reference() -> None:
|
|
50
|
+
id_token = _jwt({
|
|
51
|
+
"https://api.openai.com/profile": {"email": "dev@example.com"},
|
|
52
|
+
"https://api.openai.com/auth": {"chatgpt_account_id": "acct_123", "chatgpt_plan_type": "pro"},
|
|
53
|
+
})
|
|
54
|
+
seen: list[httpx.Request] = []
|
|
55
|
+
|
|
56
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
57
|
+
seen.append(request)
|
|
58
|
+
if request.url.path == "/api/accounts/deviceauth/usercode":
|
|
59
|
+
assert json.loads(request.content) == {"client_id": CODEX_CLIENT_ID}
|
|
60
|
+
return httpx.Response(200, json={"device_auth_id": "dev_1", "user_code": "ABCD", "interval": "1"})
|
|
61
|
+
if request.url.path == "/api/accounts/deviceauth/token":
|
|
62
|
+
assert json.loads(request.content) == {"device_auth_id": "dev_1", "user_code": "ABCD"}
|
|
63
|
+
return httpx.Response(
|
|
64
|
+
200,
|
|
65
|
+
json={"authorization_code": "auth_code", "code_challenge": "challenge", "code_verifier": "verifier"},
|
|
66
|
+
)
|
|
67
|
+
if request.url.path == "/oauth/token":
|
|
68
|
+
body = dict(pair.split("=") for pair in request.content.decode().split("&"))
|
|
69
|
+
assert body["grant_type"] == "authorization_code"
|
|
70
|
+
assert body["code"] == "auth_code"
|
|
71
|
+
assert body["client_id"] == CODEX_CLIENT_ID
|
|
72
|
+
assert body["code_verifier"] == "verifier"
|
|
73
|
+
return httpx.Response(
|
|
74
|
+
200, json={"id_token": id_token, "access_token": ACCESS_TOKEN, "refresh_token": REFRESH_TOKEN}
|
|
75
|
+
)
|
|
76
|
+
return httpx.Response(404)
|
|
77
|
+
|
|
78
|
+
client = CodexOAuthClient(http_client=httpx.Client(transport=httpx.MockTransport(handler)))
|
|
79
|
+
|
|
80
|
+
device_code, record = client.login_device_code(timeout_seconds=1)
|
|
81
|
+
|
|
82
|
+
assert device_code.user_code == "ABCD"
|
|
83
|
+
assert record.tokens.access_token == ACCESS_TOKEN
|
|
84
|
+
assert record.tokens.refresh_token == REFRESH_TOKEN
|
|
85
|
+
assert record.account.email == "dev@example.com"
|
|
86
|
+
assert record.account.chatgpt_account_id == "acct_123"
|
|
87
|
+
assert [request.url.path for request in seen] == [
|
|
88
|
+
"/api/accounts/deviceauth/usercode",
|
|
89
|
+
"/api/accounts/deviceauth/token",
|
|
90
|
+
"/oauth/token",
|
|
91
|
+
]
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
def test_codex_refresh_preserves_omitted_token_fields(tmp_path) -> None:
|
|
95
|
+
id_token = _jwt({"https://api.openai.com/auth": {"chatgpt_account_id": "acct_new"}})
|
|
96
|
+
|
|
97
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
98
|
+
assert str(request.url) == CODEX_TOKEN_ENDPOINT
|
|
99
|
+
assert json.loads(request.content) == {
|
|
100
|
+
"client_id": CODEX_CLIENT_ID,
|
|
101
|
+
"grant_type": "refresh_token",
|
|
102
|
+
"refresh_token": OLD_REFRESH_TOKEN,
|
|
103
|
+
}
|
|
104
|
+
return httpx.Response(200, json={"id_token": id_token, "access_token": NEW_ACCESS_TOKEN})
|
|
105
|
+
|
|
106
|
+
store = OAuthStore(tmp_path / "auth.json")
|
|
107
|
+
client = CodexOAuthClient(store=store, http_client=httpx.Client(transport=httpx.MockTransport(handler)))
|
|
108
|
+
record = OAuthProviderRecord(
|
|
109
|
+
issuer="https://auth.openai.com",
|
|
110
|
+
client_id=CODEX_CLIENT_ID,
|
|
111
|
+
token_endpoint=CODEX_TOKEN_ENDPOINT,
|
|
112
|
+
tokens=OAuthTokens(
|
|
113
|
+
id_token=OLD_ID_TOKEN,
|
|
114
|
+
access_token=OLD_ACCESS_TOKEN,
|
|
115
|
+
refresh_token=OLD_REFRESH_TOKEN,
|
|
116
|
+
),
|
|
117
|
+
last_refresh_at=datetime.now(UTC),
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
refreshed = client.refresh_record(record)
|
|
121
|
+
|
|
122
|
+
assert refreshed.tokens.id_token == id_token
|
|
123
|
+
assert refreshed.tokens.access_token == NEW_ACCESS_TOKEN
|
|
124
|
+
assert refreshed.tokens.refresh_token == OLD_REFRESH_TOKEN
|
|
125
|
+
assert refreshed.account.chatgpt_account_id == "acct_new"
|
|
126
|
+
|
|
127
|
+
|
|
128
|
+
def test_codex_refresh_rejects_account_mismatch(tmp_path) -> None:
|
|
129
|
+
id_token = _jwt({"https://api.openai.com/auth": {"chatgpt_account_id": "acct_new"}})
|
|
130
|
+
|
|
131
|
+
def handler(request: httpx.Request) -> httpx.Response:
|
|
132
|
+
return httpx.Response(200, json={"id_token": id_token, "access_token": NEW_ACCESS_TOKEN})
|
|
133
|
+
|
|
134
|
+
store = OAuthStore(tmp_path / "auth.json")
|
|
135
|
+
client = CodexOAuthClient(store=store, http_client=httpx.Client(transport=httpx.MockTransport(handler)))
|
|
136
|
+
record = OAuthProviderRecord(
|
|
137
|
+
issuer="https://auth.openai.com",
|
|
138
|
+
client_id=CODEX_CLIENT_ID,
|
|
139
|
+
token_endpoint=CODEX_TOKEN_ENDPOINT,
|
|
140
|
+
tokens=OAuthTokens(access_token=OLD_ACCESS_TOKEN, refresh_token=OLD_REFRESH_TOKEN),
|
|
141
|
+
account=OAuthAccount(chatgpt_account_id="acct_old"),
|
|
142
|
+
)
|
|
143
|
+
|
|
144
|
+
import pytest
|
|
145
|
+
|
|
146
|
+
with pytest.raises(RuntimeError, match="different ChatGPT account"):
|
|
147
|
+
client.refresh_record(record)
|
|
148
|
+
|
|
149
|
+
|
|
150
|
+
def test_store_permissions(tmp_path) -> None:
|
|
151
|
+
store = OAuthStore(tmp_path / ".yaai" / "auth.json")
|
|
152
|
+
store.save(store.load())
|
|
153
|
+
|
|
154
|
+
assert oct(store.path.parent.stat().st_mode & 0o777) == "0o700"
|
|
155
|
+
assert oct(store.path.stat().st_mode & 0o777) == "0o600"
|
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
|
|
5
|
+
from ya_oauth.store import OAuthStore
|
|
6
|
+
from ya_oauth.types import AuthFile
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def test_load_repairs_existing_auth_file_mode(tmp_path) -> None:
|
|
10
|
+
auth_path = tmp_path / ".yaai" / "auth.json"
|
|
11
|
+
auth_path.parent.mkdir(mode=0o700)
|
|
12
|
+
auth_path.write_text(json.dumps(AuthFile().model_dump(mode="json")), encoding="utf-8")
|
|
13
|
+
auth_path.chmod(0o644)
|
|
14
|
+
|
|
15
|
+
store = OAuthStore(auth_path)
|
|
16
|
+
|
|
17
|
+
assert store.load().version == 1
|
|
18
|
+
assert auth_path.stat().st_mode & 0o777 == 0o600
|
|
@@ -0,0 +1,15 @@
|
|
|
1
|
+
"""OAuth login, refresh, storage, and CLI for YA model providers."""
|
|
2
|
+
|
|
3
|
+
from ya_oauth.codex import CODEX_PROFILE, CodexOAuthClient
|
|
4
|
+
from ya_oauth.store import OAuthStore, StoreBackedTokenSource
|
|
5
|
+
from ya_oauth.types import OAuthAccount, OAuthProviderRecord, OAuthTokens
|
|
6
|
+
|
|
7
|
+
__all__ = [
|
|
8
|
+
"CODEX_PROFILE",
|
|
9
|
+
"CodexOAuthClient",
|
|
10
|
+
"OAuthAccount",
|
|
11
|
+
"OAuthProviderRecord",
|
|
12
|
+
"OAuthStore",
|
|
13
|
+
"OAuthTokens",
|
|
14
|
+
"StoreBackedTokenSource",
|
|
15
|
+
]
|
|
@@ -0,0 +1,146 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
import stat
|
|
5
|
+
from pathlib import Path
|
|
6
|
+
|
|
7
|
+
import click
|
|
8
|
+
|
|
9
|
+
from ya_oauth.codex import CodexOAuthClient, redact_record
|
|
10
|
+
from ya_oauth.store import DEFAULT_AUTH_PATH, OAuthStore
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
@click.group()
|
|
14
|
+
def cli() -> None:
|
|
15
|
+
"""Manage OAuth credentials for YA model providers."""
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
@cli.command()
|
|
19
|
+
@click.argument("provider", type=click.Choice(["codex"]))
|
|
20
|
+
@click.option(
|
|
21
|
+
"--auth-file", type=click.Path(path_type=Path), default=None, help="Auth file path. Defaults to ~/.yaai/auth.json."
|
|
22
|
+
)
|
|
23
|
+
def login(provider: str, auth_file: Path | None) -> None:
|
|
24
|
+
"""Log in to an OAuth provider."""
|
|
25
|
+
store = OAuthStore(auth_file)
|
|
26
|
+
if provider == "codex":
|
|
27
|
+
client = CodexOAuthClient(store=store)
|
|
28
|
+
try:
|
|
29
|
+
device_code = client.request_device_code()
|
|
30
|
+
click.echo("Open this URL in your browser and sign in to ChatGPT:")
|
|
31
|
+
click.echo(device_code.verification_url)
|
|
32
|
+
click.echo("")
|
|
33
|
+
click.echo("Enter this one-time code:")
|
|
34
|
+
click.echo(device_code.user_code)
|
|
35
|
+
click.echo("")
|
|
36
|
+
click.echo("Waiting for browser authorization...")
|
|
37
|
+
token_code = client.poll_device_token(device_code)
|
|
38
|
+
record = client.exchange_device_code(token_code)
|
|
39
|
+
store.set_provider("codex", record)
|
|
40
|
+
email = record.account.email or "unknown account"
|
|
41
|
+
click.echo(f"Logged in to codex as {email}.")
|
|
42
|
+
finally:
|
|
43
|
+
client.close()
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
@cli.command()
|
|
47
|
+
@click.argument("provider", type=click.Choice(["codex"]), required=False)
|
|
48
|
+
@click.option(
|
|
49
|
+
"--auth-file", type=click.Path(path_type=Path), default=None, help="Auth file path. Defaults to ~/.yaai/auth.json."
|
|
50
|
+
)
|
|
51
|
+
def status(provider: str | None, auth_file: Path | None) -> None:
|
|
52
|
+
"""Show OAuth provider login status."""
|
|
53
|
+
store = OAuthStore(auth_file)
|
|
54
|
+
auth = store.load()
|
|
55
|
+
provider_names = [provider] if provider else sorted(auth.providers)
|
|
56
|
+
if not provider_names:
|
|
57
|
+
click.echo("No OAuth providers are logged in.")
|
|
58
|
+
return
|
|
59
|
+
for provider_name in provider_names:
|
|
60
|
+
record = auth.providers.get(provider_name)
|
|
61
|
+
if record is None:
|
|
62
|
+
click.echo(f"{provider_name}: not logged in")
|
|
63
|
+
continue
|
|
64
|
+
account = record.account
|
|
65
|
+
identity = account.email or account.chatgpt_user_id or "unknown account"
|
|
66
|
+
plan = f", plan={account.chatgpt_plan_type}" if account.chatgpt_plan_type else ""
|
|
67
|
+
refreshed = record.last_refresh_at.isoformat() if record.last_refresh_at else "never"
|
|
68
|
+
click.echo(f"{provider_name}: logged in as {identity}{plan}, last_refresh_at={refreshed}")
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
@cli.command(name="refresh")
|
|
72
|
+
@click.argument("provider", type=click.Choice(["codex"]))
|
|
73
|
+
@click.option(
|
|
74
|
+
"--auth-file", type=click.Path(path_type=Path), default=None, help="Auth file path. Defaults to ~/.yaai/auth.json."
|
|
75
|
+
)
|
|
76
|
+
def refresh_cmd(provider: str, auth_file: Path | None) -> None:
|
|
77
|
+
"""Refresh OAuth credentials."""
|
|
78
|
+
store = OAuthStore(auth_file)
|
|
79
|
+
if provider == "codex":
|
|
80
|
+
client = CodexOAuthClient(store=store)
|
|
81
|
+
try:
|
|
82
|
+
source = client.make_token_source()
|
|
83
|
+
snapshot = _run_sync(source.refresh_token())
|
|
84
|
+
identity = snapshot.account.email or snapshot.account.chatgpt_user_id or "unknown account"
|
|
85
|
+
click.echo(f"Refreshed codex credentials for {identity}.")
|
|
86
|
+
finally:
|
|
87
|
+
client.close()
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
@cli.command()
|
|
91
|
+
@click.argument("provider", type=click.Choice(["codex"]))
|
|
92
|
+
@click.option(
|
|
93
|
+
"--auth-file", type=click.Path(path_type=Path), default=None, help="Auth file path. Defaults to ~/.yaai/auth.json."
|
|
94
|
+
)
|
|
95
|
+
@click.option("--revoke/--no-revoke", default=True, help="Revoke provider tokens before deleting local credentials.")
|
|
96
|
+
def logout(provider: str, auth_file: Path | None, revoke: bool) -> None:
|
|
97
|
+
"""Log out from an OAuth provider."""
|
|
98
|
+
store = OAuthStore(auth_file)
|
|
99
|
+
record = store.get_provider(provider)
|
|
100
|
+
if record is None:
|
|
101
|
+
click.echo(f"{provider}: not logged in")
|
|
102
|
+
return
|
|
103
|
+
if provider == "codex" and revoke:
|
|
104
|
+
client = CodexOAuthClient(store=store)
|
|
105
|
+
try:
|
|
106
|
+
client.revoke_record(record)
|
|
107
|
+
finally:
|
|
108
|
+
client.close()
|
|
109
|
+
store.delete_provider(provider)
|
|
110
|
+
click.echo(f"Logged out from {provider}.")
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
@cli.command()
|
|
114
|
+
@click.option(
|
|
115
|
+
"--auth-file", type=click.Path(path_type=Path), default=None, help="Auth file path. Defaults to ~/.yaai/auth.json."
|
|
116
|
+
)
|
|
117
|
+
def doctor(auth_file: Path | None) -> None:
|
|
118
|
+
"""Inspect OAuth store health without printing tokens."""
|
|
119
|
+
path = (auth_file or DEFAULT_AUTH_PATH).expanduser()
|
|
120
|
+
store = OAuthStore(path)
|
|
121
|
+
auth = store.load()
|
|
122
|
+
click.echo(f"Auth file: {path}")
|
|
123
|
+
click.echo(f"Providers: {', '.join(sorted(auth.providers)) if auth.providers else 'none'}")
|
|
124
|
+
_print_mode("Directory", path.parent, expected=0o700)
|
|
125
|
+
if path.exists():
|
|
126
|
+
_print_mode("File", path, expected=0o600)
|
|
127
|
+
for provider_name, record in sorted(auth.providers.items()):
|
|
128
|
+
safe_record = redact_record(record)
|
|
129
|
+
account = safe_record.get("account", {})
|
|
130
|
+
click.echo(f"{provider_name}: account={account}")
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
def _print_mode(label: str, path: Path, *, expected: int) -> None:
|
|
134
|
+
mode = stat.S_IMODE(os.stat(path).st_mode)
|
|
135
|
+
status_text = "ok" if mode == expected else f"expected {expected:o}"
|
|
136
|
+
click.echo(f"{label} mode: {mode:o} ({status_text})")
|
|
137
|
+
|
|
138
|
+
|
|
139
|
+
def _run_sync(awaitable): # type: ignore[no-untyped-def]
|
|
140
|
+
import asyncio
|
|
141
|
+
|
|
142
|
+
try:
|
|
143
|
+
loop = asyncio.get_running_loop()
|
|
144
|
+
except RuntimeError:
|
|
145
|
+
return asyncio.run(awaitable)
|
|
146
|
+
raise RuntimeError(f"Cannot run ya-oauth CLI command inside active event loop: {loop}")
|
|
@@ -0,0 +1,239 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
from datetime import UTC, datetime
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
import httpx
|
|
8
|
+
from pydantic import AliasChoices, BaseModel, Field
|
|
9
|
+
|
|
10
|
+
from ya_oauth.jwt import account_from_id_token
|
|
11
|
+
from ya_oauth.store import OAuthStore, StoreBackedTokenSource
|
|
12
|
+
from ya_oauth.types import OAuthAccount, OAuthProviderRecord, OAuthTokens
|
|
13
|
+
|
|
14
|
+
CODEX_ISSUER = "https://auth.openai.com"
|
|
15
|
+
CODEX_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
|
|
16
|
+
CODEX_TOKEN_ENDPOINT = "https://auth.openai.com/oauth/token" # noqa: S105
|
|
17
|
+
CODEX_REVOKE_ENDPOINT = "https://auth.openai.com/oauth/revoke"
|
|
18
|
+
CODEX_BASE_URL = "https://chatgpt.com/backend-api/codex"
|
|
19
|
+
CODEX_DEVICE_REDIRECT_URI = "https://auth.openai.com/deviceauth/callback"
|
|
20
|
+
CODEX_SCOPES = [
|
|
21
|
+
"openid",
|
|
22
|
+
"profile",
|
|
23
|
+
"email",
|
|
24
|
+
"offline_access",
|
|
25
|
+
"api.connectors.read",
|
|
26
|
+
"api.connectors.invoke",
|
|
27
|
+
]
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
@dataclass(frozen=True)
|
|
31
|
+
class CodexOAuthProfile:
|
|
32
|
+
issuer: str = CODEX_ISSUER
|
|
33
|
+
client_id: str = CODEX_CLIENT_ID
|
|
34
|
+
token_endpoint: str = CODEX_TOKEN_ENDPOINT
|
|
35
|
+
revoke_endpoint: str = CODEX_REVOKE_ENDPOINT
|
|
36
|
+
base_url: str = CODEX_BASE_URL
|
|
37
|
+
scopes: tuple[str, ...] = tuple(CODEX_SCOPES)
|
|
38
|
+
|
|
39
|
+
@property
|
|
40
|
+
def device_user_code_endpoint(self) -> str:
|
|
41
|
+
return f"{self.issuer.rstrip('/')}/api/accounts/deviceauth/usercode"
|
|
42
|
+
|
|
43
|
+
@property
|
|
44
|
+
def device_token_endpoint(self) -> str:
|
|
45
|
+
return f"{self.issuer.rstrip('/')}/api/accounts/deviceauth/token"
|
|
46
|
+
|
|
47
|
+
@property
|
|
48
|
+
def verification_url(self) -> str:
|
|
49
|
+
return f"{self.issuer.rstrip('/')}/codex/device"
|
|
50
|
+
|
|
51
|
+
@property
|
|
52
|
+
def device_redirect_uri(self) -> str:
|
|
53
|
+
return f"{self.issuer.rstrip('/')}/deviceauth/callback"
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
CODEX_PROFILE = CodexOAuthProfile()
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
class DeviceCode(BaseModel):
|
|
60
|
+
verification_url: str
|
|
61
|
+
user_code: str
|
|
62
|
+
device_auth_id: str
|
|
63
|
+
interval: int = 5
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
class _UserCodeResponse(BaseModel):
|
|
67
|
+
device_auth_id: str
|
|
68
|
+
user_code: str = Field(validation_alias=AliasChoices("user_code", "usercode"))
|
|
69
|
+
interval: int | str = 5
|
|
70
|
+
|
|
71
|
+
def to_device_code(self, profile: CodexOAuthProfile) -> DeviceCode:
|
|
72
|
+
return DeviceCode(
|
|
73
|
+
verification_url=profile.verification_url,
|
|
74
|
+
user_code=self.user_code,
|
|
75
|
+
device_auth_id=self.device_auth_id,
|
|
76
|
+
interval=int(self.interval),
|
|
77
|
+
)
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
class _DeviceTokenResponse(BaseModel):
|
|
81
|
+
authorization_code: str
|
|
82
|
+
code_challenge: str
|
|
83
|
+
code_verifier: str
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
class _TokenResponse(BaseModel):
|
|
87
|
+
id_token: str | None = None
|
|
88
|
+
access_token: str | None = None
|
|
89
|
+
refresh_token: str | None = None
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
class CodexOAuthClient:
|
|
93
|
+
"""Codex OAuth device-code login and refresh client aligned with OpenAI Codex."""
|
|
94
|
+
|
|
95
|
+
def __init__(
|
|
96
|
+
self,
|
|
97
|
+
*,
|
|
98
|
+
profile: CodexOAuthProfile = CODEX_PROFILE,
|
|
99
|
+
store: OAuthStore | None = None,
|
|
100
|
+
http_client: httpx.Client | None = None,
|
|
101
|
+
) -> None:
|
|
102
|
+
self.profile = profile
|
|
103
|
+
self.store = store or OAuthStore()
|
|
104
|
+
self.http_client = http_client or httpx.Client(timeout=30)
|
|
105
|
+
self._owns_http_client = http_client is None
|
|
106
|
+
|
|
107
|
+
def close(self) -> None:
|
|
108
|
+
if self._owns_http_client:
|
|
109
|
+
self.http_client.close()
|
|
110
|
+
|
|
111
|
+
def request_device_code(self) -> DeviceCode:
|
|
112
|
+
response = self.http_client.post(
|
|
113
|
+
self.profile.device_user_code_endpoint,
|
|
114
|
+
json={"client_id": self.profile.client_id},
|
|
115
|
+
headers={"Content-Type": "application/json"},
|
|
116
|
+
)
|
|
117
|
+
response.raise_for_status()
|
|
118
|
+
return _UserCodeResponse.model_validate(response.json()).to_device_code(self.profile)
|
|
119
|
+
|
|
120
|
+
def poll_device_token(self, device_code: DeviceCode, *, timeout_seconds: int = 15 * 60) -> _DeviceTokenResponse:
|
|
121
|
+
monotonic = __import__("time").monotonic
|
|
122
|
+
end_at = monotonic() + timeout_seconds
|
|
123
|
+
while True:
|
|
124
|
+
response = self.http_client.post(
|
|
125
|
+
self.profile.device_token_endpoint,
|
|
126
|
+
json={"device_auth_id": device_code.device_auth_id, "user_code": device_code.user_code},
|
|
127
|
+
headers={"Content-Type": "application/json"},
|
|
128
|
+
)
|
|
129
|
+
if response.is_success:
|
|
130
|
+
return _DeviceTokenResponse.model_validate(response.json())
|
|
131
|
+
if response.status_code in (403, 404) and monotonic() < end_at:
|
|
132
|
+
sleep_for = min(device_code.interval, max(0.0, end_at - monotonic()))
|
|
133
|
+
__import__("time").sleep(sleep_for)
|
|
134
|
+
continue
|
|
135
|
+
response.raise_for_status()
|
|
136
|
+
raise RuntimeError("Codex device authorization failed")
|
|
137
|
+
|
|
138
|
+
def exchange_device_code(self, code_response: _DeviceTokenResponse) -> OAuthProviderRecord:
|
|
139
|
+
response = self.http_client.post(
|
|
140
|
+
self.profile.token_endpoint,
|
|
141
|
+
data={
|
|
142
|
+
"grant_type": "authorization_code",
|
|
143
|
+
"code": code_response.authorization_code,
|
|
144
|
+
"redirect_uri": self.profile.device_redirect_uri,
|
|
145
|
+
"client_id": self.profile.client_id,
|
|
146
|
+
"code_verifier": code_response.code_verifier,
|
|
147
|
+
},
|
|
148
|
+
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
|
149
|
+
)
|
|
150
|
+
response.raise_for_status()
|
|
151
|
+
token_response = _TokenResponse.model_validate(response.json())
|
|
152
|
+
return self._record_from_token_response(token_response)
|
|
153
|
+
|
|
154
|
+
def login_device_code(self, *, timeout_seconds: int = 15 * 60) -> tuple[DeviceCode, OAuthProviderRecord]:
|
|
155
|
+
device_code = self.request_device_code()
|
|
156
|
+
token_code = self.poll_device_token(device_code, timeout_seconds=timeout_seconds)
|
|
157
|
+
return device_code, self.exchange_device_code(token_code)
|
|
158
|
+
|
|
159
|
+
def refresh_record(self, record: OAuthProviderRecord) -> OAuthProviderRecord:
|
|
160
|
+
refresh_token = record.tokens.refresh_token
|
|
161
|
+
if not refresh_token:
|
|
162
|
+
raise RuntimeError("Codex refresh token is missing; run `ya-oauth login codex` again.")
|
|
163
|
+
response = self.http_client.post(
|
|
164
|
+
self.profile.token_endpoint,
|
|
165
|
+
json={
|
|
166
|
+
"client_id": self.profile.client_id,
|
|
167
|
+
"grant_type": "refresh_token",
|
|
168
|
+
"refresh_token": refresh_token,
|
|
169
|
+
},
|
|
170
|
+
headers={"Content-Type": "application/json"},
|
|
171
|
+
)
|
|
172
|
+
response.raise_for_status()
|
|
173
|
+
token_response = _TokenResponse.model_validate(response.json())
|
|
174
|
+
account = account_from_id_token(token_response.id_token) if token_response.id_token else record.account
|
|
175
|
+
_validate_same_account(record.account, account)
|
|
176
|
+
return record.with_refreshed_tokens(
|
|
177
|
+
id_token=token_response.id_token,
|
|
178
|
+
access_token=token_response.access_token,
|
|
179
|
+
refresh_token=token_response.refresh_token,
|
|
180
|
+
account=account,
|
|
181
|
+
)
|
|
182
|
+
|
|
183
|
+
def revoke_record(self, record: OAuthProviderRecord) -> None:
|
|
184
|
+
token = record.tokens.refresh_token or record.tokens.access_token
|
|
185
|
+
if not token or not self.profile.revoke_endpoint:
|
|
186
|
+
return
|
|
187
|
+
response = self.http_client.post(
|
|
188
|
+
self.profile.revoke_endpoint,
|
|
189
|
+
data={"client_id": self.profile.client_id, "token": token},
|
|
190
|
+
)
|
|
191
|
+
if response.status_code < 400:
|
|
192
|
+
return
|
|
193
|
+
response.raise_for_status()
|
|
194
|
+
|
|
195
|
+
def make_token_source(self) -> StoreBackedTokenSource:
|
|
196
|
+
return StoreBackedTokenSource("codex", store=self.store, refresh_provider=self.refresh_record)
|
|
197
|
+
|
|
198
|
+
def _record_from_token_response(self, token_response: _TokenResponse) -> OAuthProviderRecord:
|
|
199
|
+
if not token_response.access_token:
|
|
200
|
+
raise RuntimeError("Codex token response did not include access_token")
|
|
201
|
+
account = OAuthAccount()
|
|
202
|
+
if token_response.id_token:
|
|
203
|
+
account = account_from_id_token(token_response.id_token)
|
|
204
|
+
return OAuthProviderRecord(
|
|
205
|
+
issuer=self.profile.issuer,
|
|
206
|
+
client_id=self.profile.client_id,
|
|
207
|
+
token_endpoint=self.profile.token_endpoint,
|
|
208
|
+
revoke_endpoint=self.profile.revoke_endpoint,
|
|
209
|
+
base_url=self.profile.base_url,
|
|
210
|
+
scopes=list(self.profile.scopes),
|
|
211
|
+
tokens=OAuthTokens(
|
|
212
|
+
id_token=token_response.id_token,
|
|
213
|
+
access_token=token_response.access_token,
|
|
214
|
+
refresh_token=token_response.refresh_token,
|
|
215
|
+
),
|
|
216
|
+
account=account,
|
|
217
|
+
last_refresh_at=datetime.now(UTC),
|
|
218
|
+
)
|
|
219
|
+
|
|
220
|
+
|
|
221
|
+
def _validate_same_account(old: OAuthAccount, new: OAuthAccount) -> None:
|
|
222
|
+
if old.chatgpt_account_id and new.chatgpt_account_id and old.chatgpt_account_id != new.chatgpt_account_id:
|
|
223
|
+
raise RuntimeError("Codex refresh returned a different ChatGPT account; run `ya-oauth login codex` again.")
|
|
224
|
+
if old.chatgpt_user_id and new.chatgpt_user_id and old.chatgpt_user_id != new.chatgpt_user_id:
|
|
225
|
+
raise RuntimeError("Codex refresh returned a different ChatGPT user; run `ya-oauth login codex` again.")
|
|
226
|
+
|
|
227
|
+
|
|
228
|
+
def create_codex_token_source(*, store: OAuthStore | None = None) -> StoreBackedTokenSource:
|
|
229
|
+
return CodexOAuthClient(store=store).make_token_source()
|
|
230
|
+
|
|
231
|
+
|
|
232
|
+
def redact_record(record: OAuthProviderRecord) -> dict[str, Any]:
|
|
233
|
+
data = record.model_dump(mode="json")
|
|
234
|
+
tokens = data.get("tokens")
|
|
235
|
+
if isinstance(tokens, dict):
|
|
236
|
+
for key in list(tokens):
|
|
237
|
+
if tokens[key]:
|
|
238
|
+
tokens[key] = "<redacted>"
|
|
239
|
+
return data
|
|
@@ -0,0 +1,50 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import base64
|
|
4
|
+
import json
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
from ya_oauth.types import OAuthAccount
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def decode_jwt_payload(jwt: str) -> dict[str, Any]:
|
|
11
|
+
"""Decode a JWT payload without signature validation for local metadata extraction."""
|
|
12
|
+
parts = jwt.split(".")
|
|
13
|
+
if len(parts) != 3 or not parts[1]:
|
|
14
|
+
raise ValueError("invalid JWT format")
|
|
15
|
+
payload = parts[1]
|
|
16
|
+
payload += "=" * (-len(payload) % 4)
|
|
17
|
+
decoded = base64.urlsafe_b64decode(payload.encode("ascii"))
|
|
18
|
+
data = json.loads(decoded.decode("utf-8"))
|
|
19
|
+
if not isinstance(data, dict):
|
|
20
|
+
raise TypeError("invalid JWT payload")
|
|
21
|
+
return data
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def account_from_id_token(id_token: str) -> OAuthAccount:
|
|
25
|
+
"""Extract Codex-compatible ChatGPT account metadata from an ID token."""
|
|
26
|
+
claims = decode_jwt_payload(id_token)
|
|
27
|
+
profile = claims.get("https://api.openai.com/profile")
|
|
28
|
+
auth = claims.get("https://api.openai.com/auth")
|
|
29
|
+
profile_data = profile if isinstance(profile, dict) else {}
|
|
30
|
+
auth_data = auth if isinstance(auth, dict) else {}
|
|
31
|
+
return OAuthAccount(
|
|
32
|
+
email=_string_or_none(claims.get("email")) or _string_or_none(profile_data.get("email")),
|
|
33
|
+
chatgpt_user_id=_string_or_none(auth_data.get("chatgpt_user_id")) or _string_or_none(auth_data.get("user_id")),
|
|
34
|
+
chatgpt_account_id=_string_or_none(auth_data.get("chatgpt_account_id")),
|
|
35
|
+
chatgpt_plan_type=_plan_type(auth_data.get("chatgpt_plan_type")),
|
|
36
|
+
chatgpt_account_is_fedramp=bool(auth_data.get("chatgpt_account_is_fedramp", False)),
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def _string_or_none(value: Any) -> str | None:
|
|
41
|
+
return value if isinstance(value, str) and value else None
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def _plan_type(value: Any) -> str | None:
|
|
45
|
+
if isinstance(value, str):
|
|
46
|
+
return value
|
|
47
|
+
if isinstance(value, dict):
|
|
48
|
+
raw = value.get("raw_value") or value.get("value") or value.get("name")
|
|
49
|
+
return _string_or_none(raw)
|
|
50
|
+
return None
|
|
@@ -0,0 +1,161 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
import contextlib
|
|
5
|
+
import fcntl
|
|
6
|
+
import json
|
|
7
|
+
import os
|
|
8
|
+
import stat
|
|
9
|
+
import tempfile
|
|
10
|
+
from collections.abc import Callable
|
|
11
|
+
from pathlib import Path
|
|
12
|
+
from typing import TypeVar
|
|
13
|
+
|
|
14
|
+
from ya_oauth.types import AuthFile, OAuthProviderRecord, TokenSnapshot
|
|
15
|
+
|
|
16
|
+
T = TypeVar("T")
|
|
17
|
+
|
|
18
|
+
DEFAULT_AUTH_DIR = Path.home() / ".yaai"
|
|
19
|
+
DEFAULT_AUTH_PATH = DEFAULT_AUTH_DIR / "auth.json"
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class OAuthStore:
|
|
23
|
+
"""File-backed OAuth credential store with process-level locking."""
|
|
24
|
+
|
|
25
|
+
def __init__(self, path: Path | str | None = None) -> None:
|
|
26
|
+
self.path = Path(path).expanduser() if path is not None else DEFAULT_AUTH_PATH
|
|
27
|
+
self.lock_path = self.path.with_suffix(self.path.suffix + ".lock")
|
|
28
|
+
|
|
29
|
+
def load(self) -> AuthFile:
|
|
30
|
+
self._ensure_parent()
|
|
31
|
+
with self._locked():
|
|
32
|
+
return self._load_unlocked()
|
|
33
|
+
|
|
34
|
+
def save(self, auth_file: AuthFile) -> None:
|
|
35
|
+
self._ensure_parent()
|
|
36
|
+
with self._locked():
|
|
37
|
+
self._save_unlocked(auth_file)
|
|
38
|
+
|
|
39
|
+
def get_provider(self, provider_name: str) -> OAuthProviderRecord | None:
|
|
40
|
+
return self.load().providers.get(provider_name)
|
|
41
|
+
|
|
42
|
+
def set_provider(self, provider_name: str, record: OAuthProviderRecord) -> None:
|
|
43
|
+
def update(auth_file: AuthFile) -> None:
|
|
44
|
+
auth_file.providers[provider_name] = record
|
|
45
|
+
|
|
46
|
+
self.update(update)
|
|
47
|
+
|
|
48
|
+
def delete_provider(self, provider_name: str) -> OAuthProviderRecord | None:
|
|
49
|
+
deleted: OAuthProviderRecord | None = None
|
|
50
|
+
|
|
51
|
+
def update(auth_file: AuthFile) -> None:
|
|
52
|
+
nonlocal deleted
|
|
53
|
+
deleted = auth_file.providers.pop(provider_name, None)
|
|
54
|
+
|
|
55
|
+
self.update(update)
|
|
56
|
+
return deleted
|
|
57
|
+
|
|
58
|
+
def update(self, updater: Callable[[AuthFile], T]) -> T:
|
|
59
|
+
self._ensure_parent()
|
|
60
|
+
with self._locked():
|
|
61
|
+
auth_file = self._load_unlocked()
|
|
62
|
+
result = updater(auth_file)
|
|
63
|
+
self._save_unlocked(auth_file)
|
|
64
|
+
return result
|
|
65
|
+
|
|
66
|
+
def _ensure_parent(self) -> None:
|
|
67
|
+
self.path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
|
|
68
|
+
with contextlib.suppress(PermissionError):
|
|
69
|
+
os.chmod(self.path.parent, 0o700)
|
|
70
|
+
|
|
71
|
+
@contextlib.contextmanager
|
|
72
|
+
def _locked(self): # type: ignore[no-untyped-def]
|
|
73
|
+
self.lock_path.parent.mkdir(mode=0o700, parents=True, exist_ok=True)
|
|
74
|
+
with self.lock_path.open("a+") as lock_file:
|
|
75
|
+
fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
|
|
76
|
+
try:
|
|
77
|
+
yield
|
|
78
|
+
finally:
|
|
79
|
+
fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
|
|
80
|
+
|
|
81
|
+
def _load_unlocked(self) -> AuthFile:
|
|
82
|
+
if not self.path.exists():
|
|
83
|
+
return AuthFile()
|
|
84
|
+
self._repair_file_mode_unlocked()
|
|
85
|
+
with self.path.open("r", encoding="utf-8") as file:
|
|
86
|
+
data = json.load(file)
|
|
87
|
+
return AuthFile.model_validate(data)
|
|
88
|
+
|
|
89
|
+
def _save_unlocked(self, auth_file: AuthFile) -> None:
|
|
90
|
+
payload = auth_file.model_dump(mode="json", exclude_none=True)
|
|
91
|
+
fd, tmp_name = tempfile.mkstemp(prefix=f".{self.path.name}.", suffix=".tmp", dir=self.path.parent, text=True)
|
|
92
|
+
tmp_path = Path(tmp_name)
|
|
93
|
+
try:
|
|
94
|
+
os.chmod(tmp_path, 0o600)
|
|
95
|
+
with os.fdopen(fd, "w", encoding="utf-8") as file:
|
|
96
|
+
json.dump(payload, file, indent=2, sort_keys=True)
|
|
97
|
+
file.write("\n")
|
|
98
|
+
file.flush()
|
|
99
|
+
os.fsync(file.fileno())
|
|
100
|
+
os.replace(tmp_path, self.path)
|
|
101
|
+
except BaseException:
|
|
102
|
+
with contextlib.suppress(FileNotFoundError):
|
|
103
|
+
tmp_path.unlink()
|
|
104
|
+
raise
|
|
105
|
+
with contextlib.suppress(PermissionError):
|
|
106
|
+
os.chmod(self.path, 0o600)
|
|
107
|
+
|
|
108
|
+
def _repair_file_mode_unlocked(self) -> None:
|
|
109
|
+
mode = stat.S_IMODE(self.path.stat().st_mode)
|
|
110
|
+
if mode != 0o600:
|
|
111
|
+
os.chmod(self.path, 0o600)
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
class StoreBackedTokenSource:
|
|
115
|
+
"""OAuthTokenSource backed by OAuthStore and a provider-specific refresher."""
|
|
116
|
+
|
|
117
|
+
def __init__(
|
|
118
|
+
self,
|
|
119
|
+
provider_name: str,
|
|
120
|
+
*,
|
|
121
|
+
store: OAuthStore | None = None,
|
|
122
|
+
refresh_provider: Callable[[OAuthProviderRecord], OAuthProviderRecord],
|
|
123
|
+
) -> None:
|
|
124
|
+
self.provider_name = provider_name
|
|
125
|
+
self.store = store or OAuthStore()
|
|
126
|
+
self._refresh_provider = refresh_provider
|
|
127
|
+
|
|
128
|
+
async def get_token(self) -> TokenSnapshot:
|
|
129
|
+
record = self.store.get_provider(self.provider_name)
|
|
130
|
+
if record is None:
|
|
131
|
+
raise RuntimeError(f"OAuth provider is not logged in: {self.provider_name}")
|
|
132
|
+
return _snapshot(self.provider_name, record)
|
|
133
|
+
|
|
134
|
+
async def refresh_token(self) -> TokenSnapshot:
|
|
135
|
+
refreshed = await asyncio.to_thread(self._refresh_and_store)
|
|
136
|
+
return _snapshot(self.provider_name, refreshed)
|
|
137
|
+
|
|
138
|
+
def _refresh_and_store(self) -> OAuthProviderRecord:
|
|
139
|
+
refreshed: OAuthProviderRecord | None = None
|
|
140
|
+
|
|
141
|
+
def update(auth_file: AuthFile) -> None:
|
|
142
|
+
nonlocal refreshed
|
|
143
|
+
record = auth_file.providers.get(self.provider_name)
|
|
144
|
+
if record is None:
|
|
145
|
+
raise RuntimeError(f"OAuth provider is not logged in: {self.provider_name}")
|
|
146
|
+
refreshed = self._refresh_provider(record)
|
|
147
|
+
auth_file.providers[self.provider_name] = refreshed
|
|
148
|
+
|
|
149
|
+
self.store.update(update)
|
|
150
|
+
if refreshed is None:
|
|
151
|
+
raise RuntimeError(f"OAuth provider refresh failed: {self.provider_name}")
|
|
152
|
+
return refreshed
|
|
153
|
+
|
|
154
|
+
|
|
155
|
+
def _snapshot(provider_name: str, record: OAuthProviderRecord) -> TokenSnapshot:
|
|
156
|
+
return TokenSnapshot(
|
|
157
|
+
provider_name=provider_name,
|
|
158
|
+
access_token=record.tokens.access_token,
|
|
159
|
+
account=record.account,
|
|
160
|
+
base_url=record.base_url,
|
|
161
|
+
)
|
|
@@ -0,0 +1,87 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from datetime import UTC, datetime
|
|
4
|
+
from typing import Any, Protocol
|
|
5
|
+
|
|
6
|
+
from pydantic import BaseModel, Field
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class OAuthTokens(BaseModel):
|
|
10
|
+
"""OAuth token material stored for a provider."""
|
|
11
|
+
|
|
12
|
+
id_token: str | None = None
|
|
13
|
+
access_token: str
|
|
14
|
+
refresh_token: str | None = None
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class OAuthAccount(BaseModel):
|
|
18
|
+
"""Account metadata derived from provider tokens."""
|
|
19
|
+
|
|
20
|
+
email: str | None = None
|
|
21
|
+
chatgpt_user_id: str | None = None
|
|
22
|
+
chatgpt_account_id: str | None = None
|
|
23
|
+
chatgpt_plan_type: str | None = None
|
|
24
|
+
chatgpt_account_is_fedramp: bool = False
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
class OAuthProviderRecord(BaseModel):
|
|
28
|
+
"""Stored OAuth configuration and credential record for one provider."""
|
|
29
|
+
|
|
30
|
+
type: str = "oauth2"
|
|
31
|
+
issuer: str
|
|
32
|
+
client_id: str
|
|
33
|
+
token_endpoint: str
|
|
34
|
+
revoke_endpoint: str | None = None
|
|
35
|
+
base_url: str | None = None
|
|
36
|
+
scopes: list[str] = Field(default_factory=list)
|
|
37
|
+
tokens: OAuthTokens
|
|
38
|
+
account: OAuthAccount = Field(default_factory=OAuthAccount)
|
|
39
|
+
last_refresh_at: datetime | None = None
|
|
40
|
+
|
|
41
|
+
def with_refreshed_tokens(
|
|
42
|
+
self,
|
|
43
|
+
*,
|
|
44
|
+
id_token: str | None = None,
|
|
45
|
+
access_token: str | None = None,
|
|
46
|
+
refresh_token: str | None = None,
|
|
47
|
+
account: OAuthAccount | None = None,
|
|
48
|
+
) -> OAuthProviderRecord:
|
|
49
|
+
"""Return an updated copy, preserving refresh response fields Codex omits."""
|
|
50
|
+
return self.model_copy(
|
|
51
|
+
update={
|
|
52
|
+
"tokens": self.tokens.model_copy(
|
|
53
|
+
update={
|
|
54
|
+
"id_token": id_token if id_token is not None else self.tokens.id_token,
|
|
55
|
+
"access_token": access_token if access_token is not None else self.tokens.access_token,
|
|
56
|
+
"refresh_token": refresh_token if refresh_token is not None else self.tokens.refresh_token,
|
|
57
|
+
}
|
|
58
|
+
),
|
|
59
|
+
"account": account if account is not None else self.account,
|
|
60
|
+
"last_refresh_at": datetime.now(UTC),
|
|
61
|
+
}
|
|
62
|
+
)
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
class AuthFile(BaseModel):
|
|
66
|
+
"""On-disk auth file schema for ~/.yaai/auth.json."""
|
|
67
|
+
|
|
68
|
+
version: int = 1
|
|
69
|
+
providers: dict[str, OAuthProviderRecord] = Field(default_factory=dict)
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
class TokenSnapshot(BaseModel):
|
|
73
|
+
"""Provider token state safe for request construction."""
|
|
74
|
+
|
|
75
|
+
provider_name: str
|
|
76
|
+
access_token: str
|
|
77
|
+
account: OAuthAccount = Field(default_factory=OAuthAccount)
|
|
78
|
+
base_url: str | None = None
|
|
79
|
+
metadata: dict[str, Any] = Field(default_factory=dict)
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
class OAuthTokenSource(Protocol):
|
|
83
|
+
"""Async token source consumed by OAuth-backed model providers."""
|
|
84
|
+
|
|
85
|
+
async def get_token(self) -> TokenSnapshot: ...
|
|
86
|
+
|
|
87
|
+
async def refresh_token(self) -> TokenSnapshot: ...
|