what._main
1import click 2 3import os.path 4 5from what.utils.file import get_file 6 7from what.cli.model import * 8from what.cli.attack import * 9from what.cli.example import * 10 11# Main CLI (what) 12@click.group() 13def main_cli(): 14 """The CLI tool for WHite-box Adversarial Toolbox (WHAT).""" 15 pass 16 17 18# what model 19@click.group() 20def model(): 21 """Manage Deep Learning Models""" 22 pass 23 24# what model list 25@model.command('list') 26def model_list(): 27 """List supported models""" 28 max_len = max([len(x[0]) for x in what_model_list]) 29 30 print(' {:<{w}s}\t{}\t{}'.format(" Model ", " Model Type ", "Description", w=max_len)) 31 print('-' * 100) 32 33 for i, model in enumerate(what_model_list, start=1): 34 if os.path.isfile(os.path.join(WHAT_MODEL_PATH, model[WHAT_MODEL_FILE_INDEX])): 35 downloaded = 'x' 36 else: 37 downloaded = ' ' 38 print('[{}] {} : {:<{w}s}\t{}\t{}'.format(downloaded, i, model[WHAT_MODEL_NAME_INDEX], model[WHAT_MODEL_TYPE_INDEX], model[WHAT_MODEL_DESC_INDEX], w=max_len)) 39 40# what model download 41@model.command('download') 42def model_download(): 43 """Download pre-trained models""" 44 max_len = max([len(x[0]) for x in what_model_list]) 45 46 print(' {:<{w}s}\t{}\t{}'.format(" Model ", " Model Type ", "Description", w=max_len)) 47 print('-' * 100) 48 49 for i, model in enumerate(what_model_list, start=1): 50 if os.path.isfile(os.path.join(WHAT_MODEL_PATH, model[WHAT_MODEL_FILE_INDEX])): 51 downloaded = 'x' 52 else: 53 downloaded = ' ' 54 print('[{}] {} : {:<{w}s}\t{}\t{}'.format(downloaded, i, model[WHAT_MODEL_NAME_INDEX], model[WHAT_MODEL_TYPE_INDEX], model[WHAT_MODEL_DESC_INDEX], w=max_len)) 55 56 index = input(f"\nPlease input the model index: ") 57 while not index.isdigit() or int(index) > len(what_model_list): 58 index = input(f"Model [{index}] does not exist. Please try again: ") 59 60 index = int(index) - 1 61 get_file(what_model_list[index][WHAT_MODEL_FILE_INDEX], 62 WHAT_MODEL_PATH, 63 what_model_list[index][WHAT_MODEL_URL_INDEX], 64 what_model_list[index][WHAT_MODEL_HASH_INDEX] 65 ) 66 67 68# what attack 69@click.group() 70def attack(): 71 """Manage Attacks""" 72 pass 73 74# what attack list 75@attack.command('list') 76def attack_list(): 77 """List supported Attacks""" 78 max_len = max([len(x[0]) for x in what_attack_list]) 79 80 for i, attack in enumerate(what_attack_list, start=1): 81 print('{} : {:<{w}s}\t{}'.format(i, attack[WHAT_ATTACK_NAME_INDEX], attack[WHAT_ATTACK_TARGET_INDEX], w=max_len)) 82 83 84# what example 85@click.group() 86def example(): 87 """Manage Examples""" 88 pass 89 90# what example list 91@example.command('list') 92def example_list(): 93 """List examples""" 94 max_len = max([len(x[WHAT_EXAMPLE_NAME_INDEX]) for x in what_example_list]) 95 96 print(' {:<{w}s}\t{}\t{}'.format(" Demo ", " Type ", "Description", w=max_len)) 97 print('-' * 80) 98 99 for i, example in enumerate(what_example_list, start=1): 100 print('{} : {:<{w}s}\t{}\t{}'.format(i, example[WHAT_EXAMPLE_NAME_INDEX], example[WHAT_EXAMPLE_TYPE_INDEX], example[WHAT_EXAMPLE_DESC_INDEX], w=max_len)) 101 102# what exmaple run 103@example.command('run') 104def example_run(): 105 """Run examples""" 106 max_len = max([len(x[0]) for x in what_example_list]) 107 108 print(' {:<{w}s}\t{}\t{}'.format(" Demo ", " Type ", "Description", w=max_len)) 109 print('-' * 80) 110 111 for i, example in enumerate(what_example_list, start=1): 112 print('{} : {:<{w}s}\t{}\t{}'.format(i, example[WHAT_EXAMPLE_NAME_INDEX], example[WHAT_EXAMPLE_TYPE_INDEX], example[WHAT_EXAMPLE_DESC_INDEX], w=max_len)) 113 114 index = input(f"\nPlease input the example index: ") 115 while not index.isdigit() or int(index) > len(what_example_list): 116 index = input(f"Example [{index}] does not exist. Please try again: ") 117 118 what_example_list[int(index)-1][WHAT_EXAMPLE_FUNC_INDEX]() 119 120def main(): 121 main_cli.add_command(model) 122 main_cli.add_command(attack) 123 main_cli.add_command(example) 124 125 model.add_command(model_list) 126 model.add_command(model_download) 127 128 attack.add_command(attack_list) 129 130 example.add_command(example_list) 131 example.add_command(example_run) 132 133 return main_cli() 134 135if __name__ == "__main__": 136 137 main()
main_cli =
<Group main-cli>
The CLI tool for WHite-box Adversarial Toolbox (WHAT).
model =
<Group model>
Manage Deep Learning Models
model_list =
<Command list>
List supported models
model_download =
<Command download>
Download pre-trained models
attack =
<Group attack>
Manage Attacks
attack_list =
<Command list>
List supported Attacks
example =
<Group example>
Manage Examples
example_list =
<Command list>
List examples
example_run =
<Command run>
Run examples
def
main():
121def main(): 122 main_cli.add_command(model) 123 main_cli.add_command(attack) 124 main_cli.add_command(example) 125 126 model.add_command(model_list) 127 model.add_command(model_download) 128 129 attack.add_command(attack_list) 130 131 example.add_command(example_list) 132 example.add_command(example_run) 133 134 return main_cli()