1import functools
2import typing
3
4import boto3
5import hyperlink
6
7
8ACCOUNT_NAMES = {
9    "760097843905": "platform",
10    "299497370133": "workflow",
11    "975596993436": "storage",
12}
13
14
15@functools.cache
16def get_aws_session(*, role_arn):
17    sts_client = boto3.client("sts")
18    assumed_role_object = sts_client.assume_role(
19        RoleArn=role_arn, RoleSessionName="AssumeRoleSession1"
20    )
21    credentials = assumed_role_object["Credentials"]
22
23    return boto3.Session(
24        aws_access_key_id=credentials["AccessKeyId"],
25        aws_secret_access_key=credentials["SecretAccessKey"],
26        aws_session_token=credentials["SessionToken"],
27    )
28
29
30def guess_account(s3_identifier, role_name):
31    """
32    Given the name of an S3 bucket, guess the account it belongs to.
33
34    You can pass the name of the bucket, or the S3 URI.
35
36    Examples:
37
38        > guess_account('s3://example-bucket/cat.jpg')
39        {'account_id': '1234567890', 'name': 'example'}
40
41        > guess_account('example-bucket')
42        {'account_id': '1234567890', 'name': 'example'}
43
44    """
45    if "wellcomedigitalworkflow" in s3_identifier:
46        account_id = "299497370133"
47    elif "wellcomecollection-storage" in s3_identifier:
48        account_id = "975596993436"
49    elif (
50        "wellcomecollection-assets-workingstorage" in s3_identifier
51        or "wellcomecollection-platform" in s3_identifier
52        or "wellcomecollection-editorial-photography" in s3_identifier
53    ):
54        account_id = "760097843905"
55    else:
56        return None
57
58    account_name = ACCOUNT_NAMES[account_id]
59
60    return {
61        "account_id": account_id,
62        "name": account_name,
63        "role_arn": f"arn:aws:iam::{account_id}:role/{account_name}-{role_name}",
64    }
65
66
67def create_s3_session(s3_identifier, *, role_name="read_only"):
68    account = guess_account(s3_identifier, role_name)
69    if account:
70        return get_aws_session(role_arn=account["role_arn"])
71    else:
72        return boto3.Session()
73
74
75class S3Uri(typing.TypedDict):
76    Bucket: str
77    Path: str
78
79
80def parse_s3_uri(s3_uri: str) -> S3Uri:
81    uri = hyperlink.parse(s3_uri)
82
83    if uri.scheme != "s3":
84        raise ValueError(f"Unrecognised scheme in {s3_uri!r}, expected s3://")
85
86    bucket = uri.host
87    path = "/".join(uri.path)
88
89    return {"Bucket": bucket, "Path": path}
90
91
92def create_link_text(*, url, label):
93    # Based on https://stackoverflow.com/a/71309268/1558022
94
95    # OSC 8 ; params ; URI ST <name> OSC 8 ;; ST
96    return f"\033]8;;{url}\033\\{label}\033]8;;\033\\"