1#!/usr/bin/env python3
2"""
3This is a rudimentary script for restoring S3 objects from Glacier.
4
5You pass it a text file with a list of S3 URIs to restore, and it will
6initiate a Glacier restore for each of them.
7
8You can also use it to track the progress of a restore operation -- it
9reports a count of how many objects are in-progress/already restored.
10"""
11
12import os
13import sys
14
15from botocore.exceptions import ClientError
16import hyperlink
17import tqdm
18
19from _common import create_s3_session
20
21sys.path.append(os.path.join(os.environ["HOME"], "repos", "concurrently"))
22from concurrently import concurrently  # noqa: E402
23
24
25def restore_object(s3_client, s3_uri):
26    uri = hyperlink.URL.from_text(s3_uri)
27
28    bucket = uri.host
29    key = "/".join(uri.path)
30
31    head_resp = s3_client.head_object(Bucket=bucket, Key=key)
32
33    if head_resp.get("Restore") == 'ongoing-request="true"':
34        return "RestoreInProgress"
35
36    if 'ongoing-request="false"' in head_resp.get("Restore", ""):
37        return "RestoredSuccessfully"
38
39    try:
40        resp = s3_client.restore_object(
41            Bucket=bucket,
42            Key=key,
43            RestoreRequest={"Days": 7, "GlacierJobParameters": {"Tier": "Standard"}},
44        )
45    except ClientError as err:
46        if err.response["Error"]["Code"] == "RestoreAlreadyInProgress":
47            return "RestoreInProgress"
48        else:
49            raise
50
51    if resp["ResponseMetadata"]["HTTPStatusCode"] == 200:
52        return "RestoredSuccessfully"
53    else:
54        return "RestoreInProgress"
55
56
57if __name__ == "__main__":
58    try:
59        path = sys.argv[1]
60    except IndexError:
61        sys.exit(f"Usage: {__file__} <LIST_OF_KEYS>")
62
63    results = {
64        "RestoredSuccessfully": 0,
65        "RestoreInProgress": 0,
66    }
67
68    with open(path) as infile:
69        s3_uris = [line.strip() for line in infile]
70
71    s3 = create_s3_session(s3_uris[0], role_name="developer").client("s3")
72
73    for _, output in tqdm.tqdm(
74        concurrently(inputs=s3_uris, handler=lambda s3_uri: restore_object(s3, s3_uri)),
75        total=len(s3_uris),
76    ):
77        results[output] += 1
78
79    from pprint import pprint
80
81    pprint(results)