bedrock-env 0.1.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.
- bedrock_env-0.1.0/LICENSE +21 -0
- bedrock_env-0.1.0/PKG-INFO +116 -0
- bedrock_env-0.1.0/README.md +95 -0
- bedrock_env-0.1.0/pyproject.toml +38 -0
- bedrock_env-0.1.0/src/bedrock_env/__init__.py +3 -0
- bedrock_env-0.1.0/src/bedrock_env/cli.py +152 -0
- bedrock_env-0.1.0/src/bedrock_env/runner.py +152 -0
- bedrock_env-0.1.0/src/bedrock_env/token.py +121 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Spring_MT
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: bedrock-env
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Generate Amazon Bedrock short-term API keys from AWS SSO and keep running processes supplied with fresh ones
|
|
5
|
+
Keywords: aws,bedrock,sso,api-key,cli
|
|
6
|
+
Author: Spring_MT
|
|
7
|
+
License-Expression: MIT
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Classifier: Environment :: Console
|
|
10
|
+
Classifier: Operating System :: POSIX
|
|
11
|
+
Classifier: Operating System :: MacOS
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Topic :: Utilities
|
|
14
|
+
Requires-Dist: aws-bedrock-token-generator>=1.1.0
|
|
15
|
+
Requires-Dist: botocore>=1.43.116
|
|
16
|
+
Requires-Python: >=3.10
|
|
17
|
+
Project-URL: Homepage, https://github.com/SpringMT/bedrock-short-time-key
|
|
18
|
+
Project-URL: Repository, https://github.com/SpringMT/bedrock-short-time-key
|
|
19
|
+
Project-URL: Issues, https://github.com/SpringMT/bedrock-short-time-key/issues
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
# bedrock-env
|
|
23
|
+
|
|
24
|
+
Generate [Amazon Bedrock short-term API keys](https://docs.aws.amazon.com/bedrock/latest/userguide/api-keys.html) from your AWS SSO profile, and keep long-running processes (dev servers, agents, ...) supplied with a fresh key.
|
|
25
|
+
|
|
26
|
+
[日本語](https://github.com/SpringMT/bedrock-short-time-key/blob/main/README.ja.md)
|
|
27
|
+
|
|
28
|
+
```sh
|
|
29
|
+
# One-off: set AWS_BEARER_TOKEN_BEDROCK in the current shell
|
|
30
|
+
eval "$(bedrock-env export --profile my-sso-profile)"
|
|
31
|
+
|
|
32
|
+
# Long-running: run a command, and restart it with a new key before the old one expires
|
|
33
|
+
bedrock-env run --profile my-sso-profile -- uv run uvicorn app:app --reload
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
## Why
|
|
37
|
+
|
|
38
|
+
A short-term Bedrock API key is a SigV4 presigned URL signed with your current AWS credentials. It is generated locally (no API call), and it stops working at whichever comes first:
|
|
39
|
+
|
|
40
|
+
- the requested lifetime (max 12 hours), or
|
|
41
|
+
- the expiry of the credentials that signed it.
|
|
42
|
+
|
|
43
|
+
With an SSO profile the second limit usually wins: role credentials last as long as the permission set's session duration (often 1 hour). Exporting a key by hand means your dev server starts failing with 403 an hour later.
|
|
44
|
+
|
|
45
|
+
A running process's environment cannot be changed from outside, so `bedrock-env run` does the next best thing: it supervises the command and restarts it with a new key shortly before the current one expires.
|
|
46
|
+
|
|
47
|
+
## Install
|
|
48
|
+
|
|
49
|
+
```sh
|
|
50
|
+
uv tool install bedrock-env
|
|
51
|
+
# or: pipx install bedrock-env
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
Requires Python 3.10+, macOS or Linux.
|
|
55
|
+
|
|
56
|
+
## Usage
|
|
57
|
+
|
|
58
|
+
### `bedrock-env export`
|
|
59
|
+
|
|
60
|
+
Prints the key so you can load it into your shell.
|
|
61
|
+
|
|
62
|
+
```sh
|
|
63
|
+
eval "$(bedrock-env export --profile my-sso-profile)" # bash / zsh
|
|
64
|
+
bedrock-env export --profile my-sso-profile --format fish | source
|
|
65
|
+
bedrock-env export --format dotenv > .env.bedrock
|
|
66
|
+
bedrock-env export --format json # {"token", "region", "expires_at"}
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
The expiry is printed to stderr.
|
|
70
|
+
|
|
71
|
+
### `bedrock-env run -- COMMAND...`
|
|
72
|
+
|
|
73
|
+
Runs `COMMAND` with `AWS_BEARER_TOKEN_BEDROCK` (and `AWS_REGION`, if unset) in its environment. Shortly before the key expires it:
|
|
74
|
+
|
|
75
|
+
1. generates a new key,
|
|
76
|
+
2. sends `SIGTERM` to the command's process group (`SIGKILL` after `--stop-timeout`),
|
|
77
|
+
3. starts the command again with the new key.
|
|
78
|
+
|
|
79
|
+
`SIGINT` / `SIGTERM` / `SIGHUP` sent to `bedrock-env` are forwarded to the command. If the command exits on its own, `bedrock-env` exits with the same status and does not restart it.
|
|
80
|
+
|
|
81
|
+
| Option | Default | |
|
|
82
|
+
| --- | --- | --- |
|
|
83
|
+
| `--profile` | `AWS_PROFILE` / default chain | AWS profile to sign with |
|
|
84
|
+
| `--region` | `AWS_REGION` / profile `region` | Region the key is valid for |
|
|
85
|
+
| `--duration` | `12` | Requested lifetime in hours (capped by the credentials' expiry) |
|
|
86
|
+
| `--margin` | `300` | Refresh this many seconds before expiry |
|
|
87
|
+
| `--retry-interval` | `60` | Wait between failed refresh attempts |
|
|
88
|
+
| `--stop-timeout` | `10` | Seconds between `SIGTERM` and `SIGKILL` on restart |
|
|
89
|
+
| `--auto-login` | off | Run `aws sso login` when the SSO session has expired |
|
|
90
|
+
|
|
91
|
+
### When the SSO session expires
|
|
92
|
+
|
|
93
|
+
The SSO session itself (`aws sso login`) also has a lifetime. Once it expires no new role credentials can be obtained:
|
|
94
|
+
|
|
95
|
+
- `bedrock-env run` keeps the current process running (its key may still be valid for a few minutes), logs `run aws sso login`, and retries every `--retry-interval`. As soon as you log in from another terminal, the next retry picks up the new session.
|
|
96
|
+
- With `--auto-login`, it runs `aws sso login --profile ...` itself, which opens a browser.
|
|
97
|
+
|
|
98
|
+
## Caveats
|
|
99
|
+
|
|
100
|
+
- **Restarts are not zero-downtime.** Requests in flight during a restart fail. With SSO this happens about once per permission set session duration.
|
|
101
|
+
- The key is bound to a single region.
|
|
102
|
+
- Only processes started by `bedrock-env run` get the new key. Shells where you ran `export` need to run it again.
|
|
103
|
+
- Windows is not supported (process groups and POSIX signals are used).
|
|
104
|
+
|
|
105
|
+
If your application uses an AWS SDK directly, you may not need an API key at all: the SDK can sign requests with your SSO profile (SigV4) and refreshes credentials itself. API keys are mainly useful for tools and SDKs that only accept a bearer token.
|
|
106
|
+
|
|
107
|
+
## Development
|
|
108
|
+
|
|
109
|
+
```sh
|
|
110
|
+
uv sync
|
|
111
|
+
uv run pytest
|
|
112
|
+
```
|
|
113
|
+
|
|
114
|
+
## License
|
|
115
|
+
|
|
116
|
+
MIT
|
|
@@ -0,0 +1,95 @@
|
|
|
1
|
+
# bedrock-env
|
|
2
|
+
|
|
3
|
+
Generate [Amazon Bedrock short-term API keys](https://docs.aws.amazon.com/bedrock/latest/userguide/api-keys.html) from your AWS SSO profile, and keep long-running processes (dev servers, agents, ...) supplied with a fresh key.
|
|
4
|
+
|
|
5
|
+
[日本語](https://github.com/SpringMT/bedrock-short-time-key/blob/main/README.ja.md)
|
|
6
|
+
|
|
7
|
+
```sh
|
|
8
|
+
# One-off: set AWS_BEARER_TOKEN_BEDROCK in the current shell
|
|
9
|
+
eval "$(bedrock-env export --profile my-sso-profile)"
|
|
10
|
+
|
|
11
|
+
# Long-running: run a command, and restart it with a new key before the old one expires
|
|
12
|
+
bedrock-env run --profile my-sso-profile -- uv run uvicorn app:app --reload
|
|
13
|
+
```
|
|
14
|
+
|
|
15
|
+
## Why
|
|
16
|
+
|
|
17
|
+
A short-term Bedrock API key is a SigV4 presigned URL signed with your current AWS credentials. It is generated locally (no API call), and it stops working at whichever comes first:
|
|
18
|
+
|
|
19
|
+
- the requested lifetime (max 12 hours), or
|
|
20
|
+
- the expiry of the credentials that signed it.
|
|
21
|
+
|
|
22
|
+
With an SSO profile the second limit usually wins: role credentials last as long as the permission set's session duration (often 1 hour). Exporting a key by hand means your dev server starts failing with 403 an hour later.
|
|
23
|
+
|
|
24
|
+
A running process's environment cannot be changed from outside, so `bedrock-env run` does the next best thing: it supervises the command and restarts it with a new key shortly before the current one expires.
|
|
25
|
+
|
|
26
|
+
## Install
|
|
27
|
+
|
|
28
|
+
```sh
|
|
29
|
+
uv tool install bedrock-env
|
|
30
|
+
# or: pipx install bedrock-env
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
Requires Python 3.10+, macOS or Linux.
|
|
34
|
+
|
|
35
|
+
## Usage
|
|
36
|
+
|
|
37
|
+
### `bedrock-env export`
|
|
38
|
+
|
|
39
|
+
Prints the key so you can load it into your shell.
|
|
40
|
+
|
|
41
|
+
```sh
|
|
42
|
+
eval "$(bedrock-env export --profile my-sso-profile)" # bash / zsh
|
|
43
|
+
bedrock-env export --profile my-sso-profile --format fish | source
|
|
44
|
+
bedrock-env export --format dotenv > .env.bedrock
|
|
45
|
+
bedrock-env export --format json # {"token", "region", "expires_at"}
|
|
46
|
+
```
|
|
47
|
+
|
|
48
|
+
The expiry is printed to stderr.
|
|
49
|
+
|
|
50
|
+
### `bedrock-env run -- COMMAND...`
|
|
51
|
+
|
|
52
|
+
Runs `COMMAND` with `AWS_BEARER_TOKEN_BEDROCK` (and `AWS_REGION`, if unset) in its environment. Shortly before the key expires it:
|
|
53
|
+
|
|
54
|
+
1. generates a new key,
|
|
55
|
+
2. sends `SIGTERM` to the command's process group (`SIGKILL` after `--stop-timeout`),
|
|
56
|
+
3. starts the command again with the new key.
|
|
57
|
+
|
|
58
|
+
`SIGINT` / `SIGTERM` / `SIGHUP` sent to `bedrock-env` are forwarded to the command. If the command exits on its own, `bedrock-env` exits with the same status and does not restart it.
|
|
59
|
+
|
|
60
|
+
| Option | Default | |
|
|
61
|
+
| --- | --- | --- |
|
|
62
|
+
| `--profile` | `AWS_PROFILE` / default chain | AWS profile to sign with |
|
|
63
|
+
| `--region` | `AWS_REGION` / profile `region` | Region the key is valid for |
|
|
64
|
+
| `--duration` | `12` | Requested lifetime in hours (capped by the credentials' expiry) |
|
|
65
|
+
| `--margin` | `300` | Refresh this many seconds before expiry |
|
|
66
|
+
| `--retry-interval` | `60` | Wait between failed refresh attempts |
|
|
67
|
+
| `--stop-timeout` | `10` | Seconds between `SIGTERM` and `SIGKILL` on restart |
|
|
68
|
+
| `--auto-login` | off | Run `aws sso login` when the SSO session has expired |
|
|
69
|
+
|
|
70
|
+
### When the SSO session expires
|
|
71
|
+
|
|
72
|
+
The SSO session itself (`aws sso login`) also has a lifetime. Once it expires no new role credentials can be obtained:
|
|
73
|
+
|
|
74
|
+
- `bedrock-env run` keeps the current process running (its key may still be valid for a few minutes), logs `run aws sso login`, and retries every `--retry-interval`. As soon as you log in from another terminal, the next retry picks up the new session.
|
|
75
|
+
- With `--auto-login`, it runs `aws sso login --profile ...` itself, which opens a browser.
|
|
76
|
+
|
|
77
|
+
## Caveats
|
|
78
|
+
|
|
79
|
+
- **Restarts are not zero-downtime.** Requests in flight during a restart fail. With SSO this happens about once per permission set session duration.
|
|
80
|
+
- The key is bound to a single region.
|
|
81
|
+
- Only processes started by `bedrock-env run` get the new key. Shells where you ran `export` need to run it again.
|
|
82
|
+
- Windows is not supported (process groups and POSIX signals are used).
|
|
83
|
+
|
|
84
|
+
If your application uses an AWS SDK directly, you may not need an API key at all: the SDK can sign requests with your SSO profile (SigV4) and refreshes credentials itself. API keys are mainly useful for tools and SDKs that only accept a bearer token.
|
|
85
|
+
|
|
86
|
+
## Development
|
|
87
|
+
|
|
88
|
+
```sh
|
|
89
|
+
uv sync
|
|
90
|
+
uv run pytest
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
## License
|
|
94
|
+
|
|
95
|
+
MIT
|
|
@@ -0,0 +1,38 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "bedrock-env"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Generate Amazon Bedrock short-term API keys from AWS SSO and keep running processes supplied with fresh ones"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
authors = [{ name = "Spring_MT" }]
|
|
7
|
+
license = "MIT"
|
|
8
|
+
license-files = ["LICENSE"]
|
|
9
|
+
requires-python = ">=3.10"
|
|
10
|
+
keywords = ["aws", "bedrock", "sso", "api-key", "cli"]
|
|
11
|
+
classifiers = [
|
|
12
|
+
"Environment :: Console",
|
|
13
|
+
"Operating System :: POSIX",
|
|
14
|
+
"Operating System :: MacOS",
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
"Topic :: Utilities",
|
|
17
|
+
]
|
|
18
|
+
dependencies = [
|
|
19
|
+
"aws-bedrock-token-generator>=1.1.0",
|
|
20
|
+
"botocore>=1.43.116",
|
|
21
|
+
]
|
|
22
|
+
|
|
23
|
+
[project.urls]
|
|
24
|
+
Homepage = "https://github.com/SpringMT/bedrock-short-time-key"
|
|
25
|
+
Repository = "https://github.com/SpringMT/bedrock-short-time-key"
|
|
26
|
+
Issues = "https://github.com/SpringMT/bedrock-short-time-key/issues"
|
|
27
|
+
|
|
28
|
+
[project.scripts]
|
|
29
|
+
bedrock-env = "bedrock_env.cli:main"
|
|
30
|
+
|
|
31
|
+
[build-system]
|
|
32
|
+
requires = ["uv_build>=0.11.28,<0.12.0"]
|
|
33
|
+
build-backend = "uv_build"
|
|
34
|
+
|
|
35
|
+
[dependency-groups]
|
|
36
|
+
dev = [
|
|
37
|
+
"pytest>=9.1.1",
|
|
38
|
+
]
|
|
@@ -0,0 +1,152 @@
|
|
|
1
|
+
"""Command line interface for bedrock-env."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import argparse
|
|
6
|
+
import json
|
|
7
|
+
import shlex
|
|
8
|
+
import subprocess
|
|
9
|
+
import sys
|
|
10
|
+
from datetime import timedelta
|
|
11
|
+
from functools import partial
|
|
12
|
+
|
|
13
|
+
from . import __version__
|
|
14
|
+
from .runner import ENV_TOKEN, Runner, log
|
|
15
|
+
from .token import LoginRequiredError, TokenError, generate_token
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def _add_aws_options(parser: argparse.ArgumentParser) -> None:
|
|
19
|
+
parser.add_argument("--profile", help="AWS profile to sign with (default: AWS_PROFILE / default chain)")
|
|
20
|
+
parser.add_argument("--region", help="Region the key is valid for (default: AWS_REGION / profile)")
|
|
21
|
+
parser.add_argument(
|
|
22
|
+
"--duration",
|
|
23
|
+
type=float,
|
|
24
|
+
default=12,
|
|
25
|
+
metavar="HOURS",
|
|
26
|
+
help="Requested lifetime in hours, max 12 (capped by the credentials' own expiry)",
|
|
27
|
+
)
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def _generator(args: argparse.Namespace):
|
|
31
|
+
return partial(
|
|
32
|
+
generate_token,
|
|
33
|
+
profile=args.profile,
|
|
34
|
+
region=args.region,
|
|
35
|
+
duration=timedelta(hours=args.duration),
|
|
36
|
+
)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def _sso_login(profile: str | None) -> None:
|
|
40
|
+
command = ["aws", "sso", "login"] + (["--profile", profile] if profile else [])
|
|
41
|
+
log(f"running `{shlex.join(command)}`")
|
|
42
|
+
subprocess.run(command, check=False)
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def cmd_export(args: argparse.Namespace) -> int:
|
|
46
|
+
token = _generator(args)()
|
|
47
|
+
expires = token.expires_at.isoformat(timespec="seconds")
|
|
48
|
+
if args.format == "json":
|
|
49
|
+
print(json.dumps({"token": token.value, "region": token.region, "expires_at": expires}))
|
|
50
|
+
elif args.format == "dotenv":
|
|
51
|
+
print(f"{ENV_TOKEN}={token.value}")
|
|
52
|
+
elif args.format == "fish":
|
|
53
|
+
print(f"set -gx {ENV_TOKEN} {shlex.quote(token.value)};")
|
|
54
|
+
else:
|
|
55
|
+
print(f"export {ENV_TOKEN}={shlex.quote(token.value)}")
|
|
56
|
+
if args.format != "json":
|
|
57
|
+
print(f"# key for {token.region}, valid until {expires}", file=sys.stderr)
|
|
58
|
+
return 0
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def cmd_run(args: argparse.Namespace) -> int:
|
|
62
|
+
command = args.command
|
|
63
|
+
if command and command[0] == "--":
|
|
64
|
+
command = command[1:]
|
|
65
|
+
if not command:
|
|
66
|
+
print("bedrock-env run: missing command (usage: bedrock-env run -- <command> ...)", file=sys.stderr)
|
|
67
|
+
return 2
|
|
68
|
+
runner = Runner(
|
|
69
|
+
command,
|
|
70
|
+
_generator(args),
|
|
71
|
+
margin=args.margin,
|
|
72
|
+
retry_interval=args.retry_interval,
|
|
73
|
+
stop_timeout=args.stop_timeout,
|
|
74
|
+
on_login_required=partial(_sso_login, args.profile) if args.auto_login else None,
|
|
75
|
+
)
|
|
76
|
+
code = runner.run()
|
|
77
|
+
# Popen reports death-by-signal as -N; shells use 128 + N.
|
|
78
|
+
return 128 - code if code < 0 else code
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def build_parser() -> argparse.ArgumentParser:
|
|
82
|
+
parser = argparse.ArgumentParser(
|
|
83
|
+
prog="bedrock-env",
|
|
84
|
+
description="Generate Amazon Bedrock short-term API keys and keep them fresh.",
|
|
85
|
+
)
|
|
86
|
+
parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
|
|
87
|
+
sub = parser.add_subparsers(dest="subcommand", required=True)
|
|
88
|
+
|
|
89
|
+
export = sub.add_parser(
|
|
90
|
+
"export",
|
|
91
|
+
help="print a key as shell exports",
|
|
92
|
+
description=f'Print {ENV_TOKEN}. Use as: eval "$(bedrock-env export)"',
|
|
93
|
+
)
|
|
94
|
+
_add_aws_options(export)
|
|
95
|
+
export.add_argument("--format", choices=["sh", "fish", "dotenv", "json"], default="sh")
|
|
96
|
+
export.set_defaults(func=cmd_export)
|
|
97
|
+
|
|
98
|
+
run = sub.add_parser(
|
|
99
|
+
"run",
|
|
100
|
+
help="run a command with a key and restart it before the key expires",
|
|
101
|
+
description=f"Run COMMAND with {ENV_TOKEN} set, regenerating the key and restarting "
|
|
102
|
+
"COMMAND shortly before it expires.",
|
|
103
|
+
)
|
|
104
|
+
_add_aws_options(run)
|
|
105
|
+
run.add_argument(
|
|
106
|
+
"--margin",
|
|
107
|
+
type=float,
|
|
108
|
+
default=300,
|
|
109
|
+
metavar="SECONDS",
|
|
110
|
+
help="refresh this many seconds before expiry (default: 300)",
|
|
111
|
+
)
|
|
112
|
+
run.add_argument(
|
|
113
|
+
"--retry-interval",
|
|
114
|
+
type=float,
|
|
115
|
+
default=60,
|
|
116
|
+
metavar="SECONDS",
|
|
117
|
+
help="wait between failed refresh attempts (default: 60)",
|
|
118
|
+
)
|
|
119
|
+
run.add_argument(
|
|
120
|
+
"--stop-timeout",
|
|
121
|
+
type=float,
|
|
122
|
+
default=10,
|
|
123
|
+
metavar="SECONDS",
|
|
124
|
+
help="wait this long after SIGTERM before SIGKILL on restart (default: 10)",
|
|
125
|
+
)
|
|
126
|
+
run.add_argument(
|
|
127
|
+
"--auto-login",
|
|
128
|
+
action="store_true",
|
|
129
|
+
help="run `aws sso login` automatically when the SSO session has expired",
|
|
130
|
+
)
|
|
131
|
+
run.add_argument("command", nargs=argparse.REMAINDER, help="command to run (after --)")
|
|
132
|
+
run.set_defaults(func=cmd_run)
|
|
133
|
+
return parser
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
def main(argv: list[str] | None = None) -> int:
|
|
137
|
+
args = build_parser().parse_args(argv)
|
|
138
|
+
try:
|
|
139
|
+
return args.func(args)
|
|
140
|
+
except LoginRequiredError as e:
|
|
141
|
+
profile = f" --profile {args.profile}" if args.profile else ""
|
|
142
|
+
print(f"bedrock-env: SSO session expired. Run `aws sso login{profile}`.\n ({e})", file=sys.stderr)
|
|
143
|
+
return 1
|
|
144
|
+
except TokenError as e:
|
|
145
|
+
print(f"bedrock-env: {e}", file=sys.stderr)
|
|
146
|
+
return 1
|
|
147
|
+
except KeyboardInterrupt:
|
|
148
|
+
return 130
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
if __name__ == "__main__":
|
|
152
|
+
sys.exit(main())
|
|
@@ -0,0 +1,152 @@
|
|
|
1
|
+
"""Run a command with a Bedrock API key and restart it before the key expires."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import os
|
|
6
|
+
import signal
|
|
7
|
+
import subprocess
|
|
8
|
+
import sys
|
|
9
|
+
import time
|
|
10
|
+
from collections.abc import Callable, Sequence
|
|
11
|
+
from datetime import datetime
|
|
12
|
+
|
|
13
|
+
from .token import LoginRequiredError, Token, TokenError
|
|
14
|
+
|
|
15
|
+
ENV_TOKEN = "AWS_BEARER_TOKEN_BEDROCK"
|
|
16
|
+
_FORWARDED_SIGNALS = (signal.SIGINT, signal.SIGTERM, signal.SIGHUP)
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def log(message: str) -> None:
|
|
20
|
+
stamp = datetime.now().strftime("%H:%M:%S")
|
|
21
|
+
print(f"[bedrock-env {stamp}] {message}", file=sys.stderr, flush=True)
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def build_env(token: Token, base: dict[str, str] | None = None) -> dict[str, str]:
|
|
25
|
+
env = dict(os.environ if base is None else base)
|
|
26
|
+
env[ENV_TOKEN] = token.value
|
|
27
|
+
# The key is bound to the region it was signed for.
|
|
28
|
+
env.setdefault("AWS_REGION", token.region)
|
|
29
|
+
return env
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
class Runner:
|
|
33
|
+
def __init__(
|
|
34
|
+
self,
|
|
35
|
+
command: Sequence[str],
|
|
36
|
+
generate: Callable[[], Token],
|
|
37
|
+
*,
|
|
38
|
+
margin: float = 300,
|
|
39
|
+
retry_interval: float = 60,
|
|
40
|
+
stop_timeout: float = 10,
|
|
41
|
+
on_login_required: Callable[[], None] | None = None,
|
|
42
|
+
):
|
|
43
|
+
self.command = list(command)
|
|
44
|
+
self.generate = generate
|
|
45
|
+
self.margin = margin
|
|
46
|
+
self.retry_interval = retry_interval
|
|
47
|
+
self.stop_timeout = stop_timeout
|
|
48
|
+
self.on_login_required = on_login_required
|
|
49
|
+
self.child: subprocess.Popen | None = None
|
|
50
|
+
self._stopping = False
|
|
51
|
+
|
|
52
|
+
# -- child process -------------------------------------------------------
|
|
53
|
+
|
|
54
|
+
def _start(self, token: Token) -> None:
|
|
55
|
+
# Own process group so restarts reach every descendant (e.g. `uv run` -> uvicorn).
|
|
56
|
+
self.child = subprocess.Popen(self.command, env=build_env(token), start_new_session=True)
|
|
57
|
+
log(f"started pid {self.child.pid}; key valid until {_local(token.expires_at)}")
|
|
58
|
+
|
|
59
|
+
def _signal_child(self, sig: int) -> None:
|
|
60
|
+
if self.child is None or self.child.poll() is not None:
|
|
61
|
+
return
|
|
62
|
+
try:
|
|
63
|
+
os.killpg(self.child.pid, sig)
|
|
64
|
+
except ProcessLookupError:
|
|
65
|
+
pass
|
|
66
|
+
|
|
67
|
+
def _stop(self) -> int:
|
|
68
|
+
assert self.child is not None
|
|
69
|
+
self._signal_child(signal.SIGTERM)
|
|
70
|
+
try:
|
|
71
|
+
return self.child.wait(timeout=self.stop_timeout)
|
|
72
|
+
except subprocess.TimeoutExpired:
|
|
73
|
+
log(f"pid {self.child.pid} did not exit in {self.stop_timeout:g}s; sending SIGKILL")
|
|
74
|
+
self._signal_child(signal.SIGKILL)
|
|
75
|
+
return self.child.wait()
|
|
76
|
+
|
|
77
|
+
def _handle_signal(self, signum, _frame) -> None:
|
|
78
|
+
self._stopping = True
|
|
79
|
+
self._signal_child(signum)
|
|
80
|
+
|
|
81
|
+
# -- main loop -----------------------------------------------------------
|
|
82
|
+
|
|
83
|
+
def _wait_child(self, seconds: float) -> int | None:
|
|
84
|
+
"""Wait up to `seconds` for the child to exit. Returns its exit code if it did."""
|
|
85
|
+
assert self.child is not None
|
|
86
|
+
deadline = time.monotonic() + max(seconds, 0)
|
|
87
|
+
while True:
|
|
88
|
+
remaining = deadline - time.monotonic()
|
|
89
|
+
try:
|
|
90
|
+
return self.child.wait(timeout=max(min(remaining, 1.0), 0))
|
|
91
|
+
except subprocess.TimeoutExpired:
|
|
92
|
+
if remaining <= 0:
|
|
93
|
+
return None
|
|
94
|
+
|
|
95
|
+
def _try_generate(self) -> Token | None:
|
|
96
|
+
try:
|
|
97
|
+
return self.generate()
|
|
98
|
+
except LoginRequiredError:
|
|
99
|
+
log("SSO session expired; run `aws sso login` (the current process keeps running)")
|
|
100
|
+
if self.on_login_required:
|
|
101
|
+
self.on_login_required()
|
|
102
|
+
try:
|
|
103
|
+
return self.generate()
|
|
104
|
+
except Exception as e:
|
|
105
|
+
log(f"failed to refresh key after login: {e}")
|
|
106
|
+
except TokenError as e:
|
|
107
|
+
log(f"failed to refresh key: {e}")
|
|
108
|
+
except Exception as e:
|
|
109
|
+
# Never let a refresh failure take the running process down with us.
|
|
110
|
+
log(f"unexpected error refreshing key: {e!r}")
|
|
111
|
+
return None
|
|
112
|
+
|
|
113
|
+
def run(self) -> int:
|
|
114
|
+
token = self.generate() # Fail fast on startup.
|
|
115
|
+
previous = {s: signal.signal(s, self._handle_signal) for s in _FORWARDED_SIGNALS}
|
|
116
|
+
try:
|
|
117
|
+
self._start(token)
|
|
118
|
+
while True:
|
|
119
|
+
# If the credentials live shorter than the margin, refresh halfway instead of
|
|
120
|
+
# restarting in a tight loop.
|
|
121
|
+
left = token.seconds_left()
|
|
122
|
+
wait = max(left - self.margin, left / 2)
|
|
123
|
+
code = self._wait_child(wait)
|
|
124
|
+
if code is not None:
|
|
125
|
+
return code
|
|
126
|
+
if self._stopping:
|
|
127
|
+
return self.child.wait()
|
|
128
|
+
|
|
129
|
+
new = self._try_generate()
|
|
130
|
+
if new is None or new.expires_at <= token.expires_at:
|
|
131
|
+
if new is not None:
|
|
132
|
+
log("credentials were not renewed yet; retrying")
|
|
133
|
+
code = self._wait_child(self.retry_interval)
|
|
134
|
+
if code is not None:
|
|
135
|
+
return code
|
|
136
|
+
continue
|
|
137
|
+
|
|
138
|
+
log("refreshing key; restarting process")
|
|
139
|
+
self._stop()
|
|
140
|
+
if self._stopping:
|
|
141
|
+
return 0
|
|
142
|
+
token = new
|
|
143
|
+
self._start(token)
|
|
144
|
+
finally:
|
|
145
|
+
if self.child is not None and self.child.poll() is None:
|
|
146
|
+
self._stop()
|
|
147
|
+
for s, handler in previous.items():
|
|
148
|
+
signal.signal(s, handler)
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
def _local(dt: datetime) -> str:
|
|
152
|
+
return dt.astimezone().strftime("%Y-%m-%d %H:%M:%S %Z")
|
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
"""Generate Amazon Bedrock short-term API keys from the AWS credential chain."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import os
|
|
6
|
+
from dataclasses import dataclass
|
|
7
|
+
from datetime import datetime, timedelta, timezone
|
|
8
|
+
|
|
9
|
+
import botocore.session
|
|
10
|
+
from aws_bedrock_token_generator import provide_token
|
|
11
|
+
from botocore import exceptions as botocore_exceptions
|
|
12
|
+
|
|
13
|
+
MAX_DURATION = timedelta(hours=12)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class TokenError(Exception):
|
|
17
|
+
"""Token could not be generated."""
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class LoginRequiredError(TokenError):
|
|
21
|
+
"""The SSO session has expired and `aws sso login` is required."""
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
# Exceptions botocore raises when the SSO access token is missing or expired.
|
|
25
|
+
_SSO_LOGIN_ERRORS = tuple(
|
|
26
|
+
getattr(botocore_exceptions, name)
|
|
27
|
+
for name in (
|
|
28
|
+
"UnauthorizedSSOTokenError",
|
|
29
|
+
"SSOTokenLoadError",
|
|
30
|
+
"TokenRetrievalError",
|
|
31
|
+
"LoginTokenLoadError",
|
|
32
|
+
)
|
|
33
|
+
if hasattr(botocore_exceptions, name)
|
|
34
|
+
)
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
@dataclass(frozen=True)
|
|
38
|
+
class Token:
|
|
39
|
+
value: str
|
|
40
|
+
region: str
|
|
41
|
+
expires_at: datetime
|
|
42
|
+
|
|
43
|
+
def seconds_left(self, now: datetime | None = None) -> float:
|
|
44
|
+
now = now or datetime.now(timezone.utc)
|
|
45
|
+
return (self.expires_at - now).total_seconds()
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
class _StaticProvider:
|
|
49
|
+
"""Credential provider that returns already-resolved credentials."""
|
|
50
|
+
|
|
51
|
+
def __init__(self, credentials):
|
|
52
|
+
self._credentials = credentials
|
|
53
|
+
|
|
54
|
+
def load(self):
|
|
55
|
+
return self._credentials
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def resolve_region(session: botocore.session.Session, region: str | None) -> str:
|
|
59
|
+
region = (
|
|
60
|
+
region
|
|
61
|
+
or os.environ.get("AWS_REGION")
|
|
62
|
+
or os.environ.get("AWS_DEFAULT_REGION")
|
|
63
|
+
or session.get_config_variable("region")
|
|
64
|
+
)
|
|
65
|
+
if not region:
|
|
66
|
+
raise TokenError(
|
|
67
|
+
"No region configured. Pass --region, set AWS_REGION, or set `region` in the profile."
|
|
68
|
+
)
|
|
69
|
+
return region
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def generate_token(
|
|
73
|
+
profile: str | None = None,
|
|
74
|
+
region: str | None = None,
|
|
75
|
+
duration: timedelta = MAX_DURATION,
|
|
76
|
+
) -> Token:
|
|
77
|
+
"""Generate a short-term Bedrock API key.
|
|
78
|
+
|
|
79
|
+
The key stops working when either `duration` passes or the credentials that signed it
|
|
80
|
+
expire, whichever comes first. For SSO profiles the latter is usually the binding limit
|
|
81
|
+
(the permission set's session duration).
|
|
82
|
+
"""
|
|
83
|
+
if not timedelta(0) < duration <= MAX_DURATION:
|
|
84
|
+
raise TokenError("Duration must be greater than 0 and at most 12 hours.")
|
|
85
|
+
|
|
86
|
+
# A fresh session on every call so botocore does not hand back cached role credentials
|
|
87
|
+
# that are about to expire.
|
|
88
|
+
try:
|
|
89
|
+
session = botocore.session.Session(profile=profile)
|
|
90
|
+
region = resolve_region(session, region)
|
|
91
|
+
credentials = session.get_credentials()
|
|
92
|
+
if credentials is None:
|
|
93
|
+
raise TokenError("No AWS credentials found.")
|
|
94
|
+
frozen = credentials.get_frozen_credentials()
|
|
95
|
+
except _SSO_LOGIN_ERRORS as e:
|
|
96
|
+
raise LoginRequiredError(str(e)) from e
|
|
97
|
+
except botocore_exceptions.ProfileNotFound as e:
|
|
98
|
+
raise TokenError(str(e)) from e
|
|
99
|
+
except (botocore_exceptions.BotoCoreError, botocore_exceptions.ClientError) as e:
|
|
100
|
+
raise TokenError(str(e)) from e
|
|
101
|
+
|
|
102
|
+
now = datetime.now(timezone.utc)
|
|
103
|
+
expires_at = now + duration
|
|
104
|
+
# RefreshableCredentials (SSO, AssumeRole, ...) know their own expiry. There is no public
|
|
105
|
+
# accessor, so read the private attribute defensively.
|
|
106
|
+
credential_expiry = getattr(credentials, "_expiry_time", None)
|
|
107
|
+
if isinstance(credential_expiry, datetime):
|
|
108
|
+
if credential_expiry.tzinfo is None:
|
|
109
|
+
credential_expiry = credential_expiry.replace(tzinfo=timezone.utc)
|
|
110
|
+
expires_at = min(expires_at, credential_expiry)
|
|
111
|
+
|
|
112
|
+
signed_for = expires_at - now
|
|
113
|
+
if signed_for.total_seconds() < 1:
|
|
114
|
+
raise TokenError("The underlying AWS credentials have already expired.")
|
|
115
|
+
|
|
116
|
+
value = provide_token(
|
|
117
|
+
region=region,
|
|
118
|
+
aws_credentials_provider=_StaticProvider(frozen),
|
|
119
|
+
expiry=timedelta(seconds=int(signed_for.total_seconds())),
|
|
120
|
+
)
|
|
121
|
+
return Token(value=value, region=region, expires_at=expires_at)
|