154 lines
5.4 KiB
Python
154 lines
5.4 KiB
Python
|
from argparse import ArgumentParser
|
||
|
from io import BytesIO
|
||
|
from json import dumps, loads
|
||
|
from ntpath import basename
|
||
|
from os import environ, listdir
|
||
|
from os.path import isfile, join
|
||
|
|
||
|
import botocore.exceptions as botoexeptions
|
||
|
from boto3 import client
|
||
|
from jinja2 import Environment, FileSystemLoader
|
||
|
|
||
|
if not environ.get("CI"):
|
||
|
from dotenv import load_dotenv
|
||
|
load_dotenv(dotenv_path="../.dev/.env")
|
||
|
|
||
|
mime = {"json": "application/json",
|
||
|
"whl": "application/x-wheel+zip", # Not a "real" type
|
||
|
"html": "text/html"
|
||
|
}
|
||
|
|
||
|
|
||
|
class Bucket():
|
||
|
def __init__(self) -> None:
|
||
|
self.client = client(
|
||
|
"s3",
|
||
|
region_name=environ.get("S3_REGION"),
|
||
|
aws_access_key_id=environ.get("S3_ACCESS_KEY"),
|
||
|
aws_secret_access_key=environ.get("S3_SECRET_KEY"),
|
||
|
endpoint_url=environ.get("S3_IS_MINIO", None)
|
||
|
)
|
||
|
self.bucket_name = environ.get("S3_BUCKET")
|
||
|
self._list()
|
||
|
|
||
|
def _list(self) -> list:
|
||
|
response = self.client.list_buckets()
|
||
|
buckets = []
|
||
|
for bucket in response["Buckets"]:
|
||
|
buckets.append(bucket)
|
||
|
return buckets
|
||
|
|
||
|
def download(self, bucket, filepath) -> bytes:
|
||
|
_file = BytesIO()
|
||
|
self.client.download_fileobj(bucket, filepath, _file)
|
||
|
_file.seek(0)
|
||
|
return _file.read()
|
||
|
|
||
|
def upload(self, bucket, filepath, file_):
|
||
|
m_type = mime[filepath.split(".")[-1]]
|
||
|
self.client.upload_fileobj(file_, bucket, filepath, ExtraArgs={
|
||
|
"ContentType": m_type, "ACL": "public-read", "CacheControl": "no-cache"})
|
||
|
|
||
|
|
||
|
class Index(Bucket):
|
||
|
def __init__(self, distribution: str, release: str, folder: str) -> None:
|
||
|
super().__init__()
|
||
|
self.distro = distribution
|
||
|
self.rel = release
|
||
|
self.path = folder
|
||
|
self.get_manifest()
|
||
|
self._gather()
|
||
|
self.loader = FileSystemLoader("templates")
|
||
|
self.jinja = Environment(loader=self.loader)
|
||
|
if environ.get("S3_IS_MINIO"):
|
||
|
self.upload_listing(self.create_listing(
|
||
|
["alpine/", "ubuntu/"]
|
||
|
), path_="index.html")
|
||
|
self.upload_listing(self.create_listing(
|
||
|
["3.13/", "3.12/"]
|
||
|
), path_="alpine/index.html")
|
||
|
self.upload_listing(self.create_listing(
|
||
|
["bionic/", "focal/"]
|
||
|
), path_="ubuntu/index.html")
|
||
|
|
||
|
def _gather(self) -> list:
|
||
|
existing_files = []
|
||
|
for key in self.manifest:
|
||
|
existing_files.append(key)
|
||
|
new_files = []
|
||
|
for file_ in listdir(self.path):
|
||
|
_dir = join(self.path, file_)
|
||
|
if isfile(_dir) and not file_.endswith("any.whl"):
|
||
|
file_name = basename(_dir)
|
||
|
if file_name not in existing_files:
|
||
|
new_files.append(file_name)
|
||
|
self.new_files = new_files
|
||
|
self.files = existing_files + new_files
|
||
|
self.files.sort(key=lambda v: v.upper())
|
||
|
|
||
|
def get_manifest(self) -> dict:
|
||
|
filepath = f"{self.distro}/{self.rel}/manifest.json"
|
||
|
try:
|
||
|
self.manifest = loads(self.download(self.bucket_name, filepath).decode("UTF-8"))
|
||
|
except botoexeptions.ClientError:
|
||
|
self.upload_manifest({})
|
||
|
self.manifest = {}
|
||
|
|
||
|
def upload_manifest(self, obj: dict):
|
||
|
filepath = f"{self.distro}/{self.rel}/manifest.json"
|
||
|
obj = bytes(dumps(obj), "ascii")
|
||
|
file_ = BytesIO(obj)
|
||
|
file_.seek(0)
|
||
|
self.upload(self.bucket_name, filepath, file_)
|
||
|
|
||
|
def add(self):
|
||
|
filepath = f"{self.distro}/{self.rel}/"
|
||
|
for wheel in self.new_files:
|
||
|
file_ = open(self.path + wheel, "rb")
|
||
|
self.upload(self.bucket_name, filepath + wheel, file_)
|
||
|
self.manifest[wheel] = {"distribution": self.distro, "release": self.rel}
|
||
|
self.upload_manifest(self.manifest)
|
||
|
self.upload_listing(self.create_listing(self.files))
|
||
|
|
||
|
def create_listing(self, wheels: list) -> str:
|
||
|
template = self.jinja.get_template("template.j2")
|
||
|
output = template.render(wheels=wheels)
|
||
|
return output
|
||
|
|
||
|
def upload_listing(self, obj: str, path_=None):
|
||
|
filepath = path_ if path_ else f"{self.distro}/{self.rel}/index.html"
|
||
|
obj = bytes(obj, "ascii")
|
||
|
file_ = BytesIO(obj)
|
||
|
file_.seek(0)
|
||
|
self.upload(self.bucket_name, filepath, file_)
|
||
|
|
||
|
|
||
|
if __name__ == "__main__":
|
||
|
parser = ArgumentParser(prog="Wheelie", description="Program to upload and present wheels to S3 buckets")
|
||
|
|
||
|
parser.add_argument("-d", "--distribution", type=str,
|
||
|
help="Distribution the wheels are target for.", choices=["alpine", "ubuntu"], required=True)
|
||
|
|
||
|
parser.add_argument("-r", "--release", type=str,
|
||
|
help="Release of the distribution.", choices=["3.12", "3.13", "focal"], required=True)
|
||
|
|
||
|
parser.add_argument("-p", "--path", type=str, default="../out/",
|
||
|
help="Path to the wheels")
|
||
|
|
||
|
parser.add_argument("-l", "--list", action="store_true", help="Lists files that would be added")
|
||
|
parser.add_argument("-a", "--add", action="store_true", help="Adds new files, and builds the index")
|
||
|
|
||
|
args = parser.parse_args()
|
||
|
|
||
|
distro = args.distribution
|
||
|
rel = args.release
|
||
|
folder = args.path
|
||
|
|
||
|
inx = Index(distribution=distro, release=rel, folder=folder)
|
||
|
|
||
|
if args.list:
|
||
|
print(inx.new_files if inx.new_files != [] else "No new files")
|
||
|
|
||
|
if args.add:
|
||
|
inx.add()
|