s3cookiejar.py

import http.cookiejar
import io

import boto3
import botocore


class S3CookieJar(http.cookiejar.LWPCookieJar):
    bucket = None
    filename = None
    encryption_key = None

    def save(self, filename=None, ignore_discard=False, ignore_expires=False):
        if filename is None:
            if self.filename is not None: filename = self.filename
            else: raise ValueError(http.cookiejar.MISSING_FILENAME_TEXT)

        session = boto3.Session()

        f = io.BytesIO()
        f.write(b"#LWP-Cookies-2.0\n")
        f.write(self.as_lwp_str(ignore_discard, ignore_expires).encode())
        f.seek(0)

        if self.encryption_key:
            kms = boto3.client("kms")
            r = kms.encrypt(KeyId=self.encryption_key, Plaintext=f.getvalue())

            f = io.BytesIO(r["CiphertextBlob"])

        s3 = session.resource("s3")
        bucket = s3.Bucket(self.bucket)
        bucket.upload_fileobj(f, filename)


    def load(self, filename=None, ignore_discard=False, ignore_expires=False):
        if filename is None:
            if self.filename is not None: filename = self.filename
            else: raise ValueError(http.cookiejar.MISSING_FILENAME_TEXT)

        f = io.BytesIO()

        session = boto3.Session()
        s3 = session.resource("s3")
        bucket = s3.Bucket(self.bucket)

        try:
            bucket.download_fileobj(filename, f)
        except botocore.exceptions.ClientError as e:
            if e.response["Error"]["Code"] == "404":
                # Don't error if the cookiejar doesn't exist yet
                return
            raise

        if self.encryption_key:
            kms = boto3.client("kms")
            r = kms.decrypt(CiphertextBlob=f.getvalue())
            f = io.BytesIO(r["Plaintext"])

        f = io.StringIO(f.getvalue().decode())

        try:
            self._really_load(f, filename, ignore_discard, ignore_expires)
        finally:
            f.close()


"""
Usage:

import requests

cookiejar = S3CookieJar()
cookiejar.bucket = "my-bucket-name"
# Optional: specify the filename inside of the bucket (includes path)
#cookiejar.filename = "some/path/to/my/cookiejar"
# Optional: specify the ARN for a KMS key to use for encryption
#cookiejar.encryption_key = ""

session = requests.Session()
session.cookies = cookiejar

cookiejar.load("test.cookiejar", ignore_discard=True, ignore_expires=True)
print(cookiejar)  # First time you run this, it'll be empty. The second time, it'll load from S3

session.get("http://httpbin.org/cookies/set?foo=bar")
print(cookiejar)  # Confirm that there's a 'foo' cookie

cookiejar.save("test.cookiejar", ignore_discard=True, ignore_expires=True)
"""