q1/main.py

import os

import boto3
from flask import Flask, redirect, render_template, send_file
from gestalt import Gestalt


def sanitize_key(key):
    return key.replace("/", "-")


def rebuild_key(key):
    return key.replace("-", "/")


def get_path(key):
    s = key.split("/")
    return "/".join(s[: len(s) - 1])


def get_filename(key):
    s = key.split("/")
    return s[len(s) - 1]


def build_s3_client(config: Gestalt):
    session = boto3.Session()
    access = config.get_string("aws.access_key")
    secret = config.get_string("aws.secret_key")
    url = config.get_string("aws.endpoint_url", "s3.amazonaws.com")
    s3 = session.client(
        service_name="s3",
        aws_access_key_id=access,
        aws_secret_access_key=secret,
        endpoint_url=url,
    )
    return s3


def list_buckets(client):
    ret = []
    list = client.list_buckets()["Buckets"]
    for bucket in list:
        ret.append(bucket["Name"])
    ret.sort()
    return ret


def list_in_bucket(client, bucket, prefix):
    ret = []
    items = []

    if prefix is None or len(prefix) == 0:
        items = client.list_objects_v2(Bucket=bucket, Delimiter="/")
    else:
        items = client.list_objects_v2(
            Bucket=bucket, Prefix=f"{prefix}/", Delimiter="/"
        )

    for item in items.get("Contents", []):
        if item.get("Key", None) is not None:
            ret.append(item["Key"])
    ret.sort()

    dirs = []
    for item in items.get("CommonPrefixes", []):
        if item.get("Prefix", None) is not None:
            dirs.append(item["Prefix"])

    dirs.sort()
    return ret, dirs


def download_object(client, bucket_name, key):
    key = rebuild_key(key)
    filename = "/tmp" + get_filename(key)
    client.download_file(Bucket=bucket_name, Key=key, Filename=filename)
    return filename


env = os.environ.get("ENV", "localkube")

app = Flask(__name__)
g = Gestalt()
g.add_config_file(f"./config/{env}.yaml")
g.build_config()
g.auto_env()

s3client = build_s3_client(g)


@app.route("/download/<bucket>/<path:path>")
def download(bucket, path):
    global s3client
    fname = download_object(s3client, bucket, path)
    return send_file(fname, as_attachment=True)


@app.route("/browse/<bucket>", defaults={"path": ""})
@app.route("/browse/<bucket>/", defaults={"path": ""})
@app.route("/browse/<bucket>/<path:path>")
@app.route("/browse/<bucket>/<path:path>/")
def withinBucket(bucket, path):
    global s3client
    items, dirs = list_in_bucket(s3client, bucket, path)
    item_names = []
    dir_names = []
    for item in items:
        item_names.append(item.split("/")[-1])
    for item in dirs:
        dir_names.append(item.split("/")[-2])
    return render_template(
        "browser.html",
        items=items,
        dirs=dirs,
        bucket=bucket,
        path=path,
        itemnames=item_names,
        dirnames=dir_names,
    )


@app.route("/browse")
@app.route("/browse/")
def sendhome():
    return redirect("/", code=302)


@app.route("/")
def root():
    global s3client
    buckets = list_buckets(s3client)
    return render_template("index.html", buckets=buckets)


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8081, debug=True)