terraform-ci-infra/1n_nmd/tools/artifacts.py

   1 #!/usr/bin/python3
   2
   3 # Copyright (c) 2020 Cisco and/or its affiliates.
   4 # Licensed under the Apache License, Version 2.0 (the "License");
   5 # you may not use this file except in compliance with the License.
   6 # You may obtain a copy of the License at:
   7 #
   8 #     http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15
  16 """Storage utilities library."""
  17
  18 import argparse
  19 import gzip
  20 import os
  21 from mimetypes import MimeTypes
  22
  23 from boto3 import resource
  24 from botocore.client import Config
  25
  26
  27 ENDPOINT_URL = u"http://storage.service.consul:9000"
  28 AWS_ACCESS_KEY_ID = u"storage"
  29 AWS_SECRET_ACCESS_KEY = u"Storage1234"
  30 REGION_NAME = u"yul1"
  31 COMPRESS_MIME = (
  32     u"text/html",
  33     u"text/xml",
  34     u"application/octet-stream"
  35 )
  36
  37
  38 def compress(src_fpath):
  39     """Compress a single file.
  40
  41     :param src_fpath: Input file path.
  42     :type src_fpath: str
  43     """
  44     with open(src_fpath, u"rb") as orig_file:
  45         with gzip.open(f"{src_fpath}.gz", u"wb") as zipped_file:
  46             zipped_file.writelines(orig_file)
  47
  48
  49 def upload(storage, bucket, src_fpath, dst_fpath):
  50     """Upload single file to destination bucket.
  51
  52     :param storage: S3 storage resource.
  53     :param bucket: S3 bucket name.
  54     :param src_fpath: Input file path.
  55     :param dst_fpath: Destination file path on remote storage.
  56     :type storage: Object
  57     :type bucket: str
  58     :type src_fpath: str
  59     :type dst_fpath: str
  60     """
  61     mime = MimeTypes().guess_type(src_fpath)[0]
  62     if not mime:
  63         mime = "application/octet-stream"
  64
  65     if mime in COMPRESS_MIME and bucket in "logs":
  66         compress(src_fpath)
  67         src_fpath = f"{src_fpath}.gz"
  68         dst_fpath = f"{dst_fpath}.gz"
  69
  70     storage.Bucket(f"{bucket}.fd.io").upload_file(
  71         src_fpath,
  72         dst_fpath,
  73         ExtraArgs={
  74             u"ContentType": mime
  75         }
  76     )
  77     print(f"https://{bucket}.nginx.service.consul/{dst_fpath}")
  78
  79
  80 def upload_recursive(storage, bucket, src_fpath):
  81     """Recursively uploads input folder to destination.
  82
  83     Example:
  84       - bucket: logs
  85       - src_fpath: /home/user
  86       - dst_fpath: logs.fd.io/home/user
  87
  88     :param storage: S3 storage resource.
  89     :param bucket: S3 bucket name.
  90     :param src_fpath: Input folder path.
  91     :type storage: Object
  92     :type bucket: str
  93     :type src_fpath: str
  94     """
  95     for path, _, files in os.walk(src_fpath):
  96         for file in files:
  97             _path = path.replace(src_fpath, u"")
  98             _dir = src_fpath[1:] if src_fpath[0] == "/" else src_fpath
  99             _dst_fpath = os.path.normpath(f"{_dir}/{_path}/{file}")
 100             _src_fpath = os.path.join(path, file)
 101             upload(storage, bucket, _src_fpath, _dst_fpath)
 102
 103
 104 def main():
 105     """Main function for storage manipulation."""
 106
 107     parser = argparse.ArgumentParser()
 108     parser.add_argument(
 109         u"-d", u"--dir", required=True, type=str,
 110         help=u"Directory to upload to storage."
 111     )
 112     parser.add_argument(
 113         u"-b", u"--bucket", required=True, type=str,
 114         help=u"Target bucket on storage."
 115     )
 116     args = parser.parse_args()
 117
 118     # Create main storage resource.
 119     storage = resource(
 120         u"s3",
 121         endpoint_url=ENDPOINT_URL,
 122         aws_access_key_id=AWS_ACCESS_KEY_ID,
 123         aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
 124         config=Config(
 125             signature_version=u"s3v4"
 126         ),
 127         region_name=REGION_NAME
 128     )
 129
 130     upload_recursive(
 131         storage=storage,
 132         bucket=args.bucket,
 133         src_fpath=args.dir
 134     )
 135
 136
 137 if __name__ == u"__main__":
 138     main()