Coverage for yaptide/admin/simulators.py: 73%
111 statements
« prev ^ index » next coverage.py v7.6.4, created at 2024-11-22 07:31 +0000
1#! /usr/bin/env python
3from functools import wraps
4import os
5from pathlib import Path
6import click
8from dotenv import load_dotenv
10if __name__ == "__main__":
11 # we need to add the main directory to the path so that we can import the yaptide package
12 import sys
13 this_file_path = Path(__file__).resolve()
14 sys.path.insert(0, str(this_file_path.parent.parent.parent))
15from yaptide.admin.simulator_storage import (decrypt_file, download_fluka_from_s3,
16 download_shieldhit_from_s3_or_from_website, download_topas_from_s3,
17 encrypt_file, upload_file_to_s3)
19load_dotenv()
20endpoint = os.getenv('S3_ENDPOINT')
21access_key = os.getenv('S3_ACCESS_KEY')
22secret_key = os.getenv('S3_SECRET_KEY')
23password = os.getenv('S3_ENCRYPTION_PASSWORD')
24salt = os.getenv('S3_ENCRYPTION_SALT')
25shieldhit_bucket = os.getenv('S3_SHIELDHIT_BUCKET')
26shieldhit_key = os.getenv('S3_SHIELDHIT_KEY')
27topas_bucket_name = os.getenv('S3_TOPAS_BUCKET')
28topas_key = os.getenv('S3_TOPAS_KEY')
29topas_version = os.getenv('S3_TOPAS_VERSION')
30geant4_bucket_name = os.getenv('S3_GEANT4_BUCKET')
31fluka_bucket = os.getenv('S3_FLUKA_BUCKET')
32fluka_key = os.getenv('S3_FLUKA_KEY')
35@click.group()
36def run():
37 """Manage simulators"""
40def s3credentials(required: bool = False):
41 """Collection of options for S3 credentials"""
43 def decorator(func):
44 """Decorator for S3 credentials options"""
46 @click.option('--endpoint',
47 type=click.STRING,
48 required=required,
49 envvar='S3_ENDPOINT',
50 default=endpoint,
51 help='S3 endpoint')
52 @click.option('--access_key',
53 type=click.STRING,
54 required=required,
55 envvar='S3_ACCESS_KEY',
56 default=access_key,
57 help='S3 access key')
58 @click.option('--secret_key',
59 type=click.STRING,
60 required=required,
61 envvar='S3_SECRET_KEY',
62 default=secret_key,
63 help='S3 secret key')
64 @wraps(func)
65 def wrapper(*args, **kwargs):
66 return func(*args, **kwargs)
68 return wrapper
70 return decorator
73def encryption_options(required: bool = False):
74 """Collection of options for encryption"""
76 def decorator(func):
77 """Decorator for encryption options"""
79 @click.option('--password',
80 type=click.STRING,
81 envvar='S3_ENCRYPTION_PASSWORD',
82 default=password,
83 required=required,
84 help='encryption password')
85 @click.option('--salt',
86 type=click.STRING,
87 envvar='S3_ENCRYPTION_SALT',
88 default=password,
89 required=required,
90 help='encryption salt')
91 @wraps(func)
92 def wrapper(*args, **kwargs):
93 return func(*args, **kwargs)
95 return wrapper
97 return decorator
100@run.command
101@click.option('--infile',
102 type=click.Path(exists=True, readable=True, file_okay=True, dir_okay=False, path_type=Path),
103 required=True,
104 help='file to encrypt')
105@click.option('--outfile',
106 type=click.Path(writable=True, file_okay=True, dir_okay=False, path_type=Path),
107 required=True,
108 help='Path where encrypted file is saved')
109@encryption_options(required=True)
110def encrypt(**kwargs):
111 """Encrypt a file"""
112 encrypted_bytes = encrypt_file(file_path=kwargs['infile'], password=kwargs['password'], salt=kwargs['salt'])
113 outfile_path = Path(kwargs['outfile'])
114 outfile_path.parent.mkdir(parents=True, exist_ok=True)
115 outfile_path.write_bytes(encrypted_bytes)
118@run.command
119@click.option('--infile',
120 type=click.Path(exists=True, readable=True, file_okay=True, dir_okay=False, path_type=Path),
121 required=True,
122 help='file to decrypt')
123@click.option('--outfile',
124 type=click.Path(writable=True, file_okay=True, dir_okay=False, path_type=Path),
125 required=True,
126 help='Path where decrypted file is saved')
127@encryption_options(required=True)
128def decrypt(**kwargs):
129 """Decrypt a file"""
130 decrypted_bytes = decrypt_file(file_path=kwargs['infile'], password=kwargs['password'], salt=kwargs['salt'])
131 outfile_path = Path(kwargs['outfile'])
132 outfile_path.parent.mkdir(parents=True, exist_ok=True)
133 outfile_path.write_bytes(decrypted_bytes)
136@run.command
137@click.option('--dir',
138 required=True,
139 type=click.Path(writable=True, file_okay=False, dir_okay=True, path_type=Path),
140 help='download directory')
141@click.option('--bucket', type=click.STRING, envvar='S3_FLUKA_BUCKET', required=True, help='S3 bucket name')
142@click.option('--key', type=click.STRING, envvar='S3_FLUKA_KEY', required=True, help='S3 key (filename)')
143@s3credentials(required=True)
144@encryption_options(required=True)
145def download_fluka(**kwargs):
146 """Download Fluka simulator"""
147 click.echo(f'Downloading Fluka into directory {kwargs["dir"]}')
148 download_status = download_fluka_from_s3(download_dir=kwargs['dir'],
149 endpoint=kwargs['endpoint'],
150 access_key=kwargs['access_key'],
151 secret_key=kwargs['secret_key'],
152 bucket=kwargs['bucket'],
153 key=kwargs['key'],
154 password=kwargs['password'],
155 salt=kwargs['salt'])
156 if download_status:
157 click.echo('Fluka downloaded')
158 else:
159 click.echo('Fluka download failed')
162@run.command
163@click.option('--dir',
164 required=True,
165 type=click.Path(writable=True, file_okay=False, dir_okay=True, path_type=Path),
166 help='download directory')
167@click.option('--topas_bucket',
168 type=click.STRING,
169 envvar='S3_TOPAS_BUCKET',
170 required=True,
171 help='S3 bucket name with TOPAS binary')
172@click.option('--geant4_bucket',
173 type=click.STRING,
174 envvar='S3_GEANT4_BUCKET',
175 required=True,
176 help='S3 bucket name with Geant4 data')
177@click.option('--topas_key', type=click.STRING, envvar='S3_TOPAS_KEY', required=True, help='S3 key (filename)')
178@click.option('--topas_version', type=click.STRING, envvar='S3_TOPAS_VERSION', required=True, help='TOPAS version')
179@s3credentials(required=True)
180def download_topas(**kwargs):
181 """Download TOPAS simulator and Geant4 data"""
182 click.echo(f'Downloading TOPAS into directory {kwargs["dir"]}')
183 installation_status = download_topas_from_s3(
184 download_dir=kwargs['dir'],
185 endpoint=kwargs['endpoint'],
186 access_key=kwargs['access_key'],
187 secret_key=kwargs['secret_key'],
188 bucket=kwargs['topas_bucket'],
189 key=kwargs['topas_key'],
190 version=kwargs['topas_version'],
191 geant4_bucket=kwargs['geant4_bucket'],
192 )
193 if installation_status:
194 click.echo('TOPAS installed')
195 else:
196 click.echo('TOPAS installation failed')
199@run.command
200@click.option('--dir',
201 required=True,
202 type=click.Path(writable=True, file_okay=False, dir_okay=True, path_type=Path),
203 help='download directory')
204@click.option('--bucket', type=click.STRING, envvar='S3_SHIELDHIT_BUCKET', help='S3 bucket name')
205@click.option('--key', type=click.STRING, envvar='S3_SHIELDHIT_KEY', help='S3 key (filename)')
206@click.option('--decrypt', is_flag=True, default=False, help='decrypt file downloaded from S3')
207@s3credentials(required=False)
208@encryption_options(required=False)
209def download_shieldhit(**kwargs):
210 """Download SHIELD-HIT12A"""
211 click.echo(f'Downloading SHIELD-HIT12A into directory {kwargs["dir"]}')
212 download_shieldhit_from_s3_or_from_website(destination_dir=kwargs['dir'],
213 endpoint=kwargs['endpoint'],
214 access_key=kwargs['access_key'],
215 secret_key=kwargs['secret_key'],
216 password=kwargs['password'],
217 salt=kwargs['salt'],
218 bucket=kwargs['bucket'],
219 key=kwargs['key'],
220 decrypt=kwargs['decrypt'])
223@run.command
224@click.option('--bucket', type=click.STRING, required=True, help='S3 bucket name')
225@click.option('--file',
226 type=click.Path(exists=True, readable=True, file_okay=True, dir_okay=False, path_type=Path),
227 required=True,
228 help='file to upload')
229@click.option('--encrypt', is_flag=True, default=False, help='encrypt file uploaded to S3')
230@s3credentials(required=True)
231@encryption_options(required=False)
232def upload(**kwargs):
233 """Upload simulator file to S3 bucket"""
234 click.echo(f'Uploading file {kwargs["file"]} to bucket {kwargs["bucket"]}')
235 upload_status = upload_file_to_s3(bucket=kwargs['bucket'],
236 file_path=kwargs['file'],
237 endpoint=kwargs['endpoint'],
238 access_key=kwargs['access_key'],
239 secret_key=kwargs['secret_key'],
240 encrypt=kwargs['encrypt'],
241 encryption_password=kwargs['password'],
242 encryption_salt=kwargs['salt'])
243 if upload_status:
244 click.echo('File uploaded successfully')
247if __name__ == "__main__":
248 run()