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.
@@ -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,3 @@
1
+ """Generate Amazon Bedrock short-term API keys and keep them fresh."""
2
+
3
+ __version__ = "0.1.0"
@@ -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)