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()