114 lines
4.3 KiB
Python
114 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import logging
|
|
import boto3
|
|
from botocore.exceptions import ClientError, ProfileNotFound
|
|
import sys
|
|
import subprocess
|
|
import re
|
|
|
|
logging.basicConfig(
|
|
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
logger = logging.getLogger("Runner")
|
|
logger.setLevel(logging.INFO)
|
|
|
|
|
|
class STS:
|
|
def __init__(self, sts_client) -> None:
|
|
self.sts = sts_client
|
|
|
|
def get_account_number(self) -> str:
|
|
try:
|
|
return self.sts.get_caller_identity()["Account"]
|
|
except Exception as e:
|
|
logger.error(f"Failed to get profile account number: {e}")
|
|
raise e
|
|
|
|
def arg_parser() -> dict:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--profile', required = False, help = "AWS profile")
|
|
parser.add_argument('--project', required = False, default = 'comp23', help = "ADVANCED USERS ONLY: Name of the project (default: 'comp23').")
|
|
parser.add_argument('--leader', required = False, help = "Docker tag name of the leader container")
|
|
parser.add_argument('--worker', required = False, help = "Docker tag name of the worker container (only for cloud track)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# argument checking.
|
|
# matcher = re.compile(allowable_project_names_mask)
|
|
if not re.match(r"^[a-z]([a-z\d\/_-])*$", args.project):
|
|
logger.error(f"Invalid project name '{args.project}'.")
|
|
logger.error(f"Project names must start with a lowercase letter and can consist only of lowercase numbers, digits, '/', '_', or '-' characters.")
|
|
sys.exit(-1)
|
|
|
|
if not (args.leader or args.worker):
|
|
logger.error(f"Neither leader or worker image specified...nothing to upload!")
|
|
sys.exit(-1)
|
|
|
|
return args
|
|
|
|
|
|
def main() -> None:
|
|
args = arg_parser()
|
|
|
|
try:
|
|
session = boto3.Session(profile_name=args.profile)
|
|
except ProfileNotFound as e:
|
|
logger.error(f"Unable to create AWS session. Profile '{profile}' not found. Please double check that this profile is set up in the ~/.aws/config file")
|
|
sys.exit(1)
|
|
if not session.region_name:
|
|
logger.error(f"Profile does not have a region defined. Please add a region (recommend: us-east-1) to profile '{profile}' in the ~/.aws/config file")
|
|
sys.exit(1)
|
|
|
|
# link to AWS services
|
|
sts_client = session.client('sts')
|
|
sts = STS(sts_client)
|
|
|
|
# in order to make the call, I need:
|
|
# - the account #
|
|
# - the region
|
|
# - the profile
|
|
# - the project id
|
|
# - the leader container
|
|
# - the worker container
|
|
|
|
account_number = sts.get_account_number()
|
|
region = session.region_name
|
|
project_name = args.project
|
|
profile_args = ['-p', args.profile] if args.profile else []
|
|
leader_args = ['-l', args.leader] if args.leader else []
|
|
worker_args = ['-w', args.worker] if args.worker else []
|
|
|
|
|
|
# run the shell script file to set it up.
|
|
cmd = ' '.join(['./ecr-push-internal.sh', \
|
|
'-j', project_name, \
|
|
'-a', str(account_number), \
|
|
'-r', region] + \
|
|
profile_args + leader_args + worker_args)
|
|
|
|
logger.info(f"About to run: { cmd }")
|
|
result = subprocess.run(cmd, shell=True)
|
|
if result.returncode == 1:
|
|
logger.error("The login command to ECR failed. Please double check that you have the correct rights in the profile.")
|
|
exit(1)
|
|
elif result.returncode == 2:
|
|
logger.error(f"The attempt to tag the leader image failed. Please make sure that you have a local leader image associated with {args.leader}.")
|
|
exit(2)
|
|
elif result.returncode == 3:
|
|
logger.error(f"The attempt to upload the leader failed. Please make sure that there is a repository associated with {project_name}.")
|
|
exit(3)
|
|
elif result.returncode == 4:
|
|
logger.error(f"The attempt to tag the worker image failed. Please make sure that you have a local worker image associated with {args.worker}.")
|
|
exit(4)
|
|
elif result.returncode == 3:
|
|
logger.error(f"The attempt to upload the worker failed. Please make sure that there is a repository associated with {project_name}.")
|
|
exit(5)
|
|
elif result.returncode == 0:
|
|
logger.info("Upload succeeded. ")
|
|
exit(0)
|
|
else:
|
|
logger.error(f"Unexpected return code: {result.returncode} from ecr-push-internal.sh; exiting")
|
|
exit(6)
|
|
|
|
if __name__ == "__main__":
|
|
main() |