| #!/usr/bin/env python |
| # Copyright 2016 The Fuchsia Authors |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| """The Cobalt command-line interface.""" |
| |
| import argparse |
| import logging |
| import os |
| import shutil |
| import subprocess |
| import sys |
| |
| import analyzers.analyzer as analyzer |
| import fake_data.generate_fake_data as fake_data |
| import randomizers.randomizer as randomizer |
| import shufflers.shuffler as shuffler |
| import tests.e2e.end_to_end_test as end_to_end_test |
| import utils.file_util as file_util |
| import visualization.generate_data_js as visualization |
| |
| THIS_DIR = os.path.dirname(__file__) |
| |
| _logger = logging.getLogger() |
| _verbose_count = 0 |
| |
| # Should public key encryption be used for communication between the |
| # Randomizers and the Analyzers via the Shufflers? |
| _use_public_key_encryption=True |
| |
| def _initLogging(verbose_count): |
| """Ensures that the logger (obtained via logging.getLogger(), as usual) is |
| initialized, with the log level set as appropriate for |verbose_count| |
| instances of --verbose on the command line.""" |
| assert(verbose_count >= 0) |
| if verbose_count == 0: |
| level = logging.WARNING |
| elif verbose_count == 1: |
| level = logging.INFO |
| else: # verbose_count >= 2 |
| level = logging.DEBUG |
| logging.basicConfig(format="%(relativeCreated).3f:%(levelname)s:%(message)s") |
| logger = logging.getLogger() |
| logger.setLevel(level) |
| logger.debug("Initialized logging: verbose_count=%d, level=%d" % |
| (verbose_count, level)) |
| logger.debug("_use_public_key_encryption=%s" % _use_public_key_encryption) |
| |
| def _build_fastrand(): |
| savedir = os.getcwd() |
| os.chdir(os.path.join(THIS_DIR, 'third_party', 'fastrand')) |
| subprocess.call(['./build.sh']) |
| os.chdir(savedir) |
| |
| def _build_fastem(): |
| savedir = os.getcwd() |
| os.chdir(os.path.join(THIS_DIR, 'third_party', 'rappor', 'analysis', |
| 'cpp')) |
| subprocess.call(['./run.sh', 'build-fast-em']) |
| os.chdir(savedir) |
| |
| def _build(): |
| _build_fastrand() |
| _build_fastem() |
| |
| def _run_end_to_end_test(): |
| _run_all() |
| end_to_end_test.check_results() |
| |
| def _test(): |
| _run_end_to_end_test() |
| |
| def _clean_all(): |
| print "Deleting the out directory..." |
| shutil.rmtree(file_util.OUT_DIR, ignore_errors=True) |
| |
| def _generate(): |
| # Generates fake data and runs the straight-counting pipeline |
| fake_data.main() |
| |
| def _randomize(): |
| # Run the randomizers |
| randomizer.readAndRandomize(_use_public_key_encryption) |
| |
| def _shuffle(): |
| # Run the shufflers |
| shuffler.main() |
| |
| def _analyze(): |
| # Run the analyzers |
| analyzer.runAllAnalyzers(_use_public_key_encryption) |
| |
| def _visualize(): |
| # Generate the visualization |
| visualization.main() |
| |
| def _run_all(): |
| _clean_all() |
| _generate() |
| _randomize() |
| _shuffle() |
| _analyze() |
| _visualize() |
| print "Done." |
| |
| def main(): |
| parser = argparse.ArgumentParser(description='The Cobalt command-line ' |
| 'interface.') |
| |
| # Note(rudominer) A note about the handling of optional arguments here. |
| # We create |parent_parser| and make it a parent of all of our sub parsers. |
| # When we want to add a global optional argument (i.e. one that applies |
| # to all sub-commands such as --verbose) we add the optional argument |
| # to both |parent_parser| and |parser|. The reason for this is that |
| # that appears to be the only way to get the help string to show up both |
| # when the top-level command is invoked and when |
| # a sub-command is invoked. |
| # |
| # In other words when the user types: |
| # |
| # python cobalt.py -h |
| # |
| # and also when the user types |
| # |
| # python cobalt.py run -h |
| # |
| # we want to show the help for the --verbose option. |
| parent_parser = argparse.ArgumentParser(add_help=False) |
| |
| parser.add_argument('--verbose', |
| help='Be verbose (multiple times for more)', |
| default=0, dest='verbose_count', action='count') |
| parent_parser.add_argument('--verbose', |
| help='Be verbose (multiple times for more)', |
| default=0, dest='verbose_count', action='count') |
| |
| |
| parser.add_argument('--no-use-encryption', |
| help='Do not us public key encryption for communication between the ' |
| 'randomizers and the analyzers via the shufflers. By default encryption ' |
| 'is used.', |
| dest="use_encryption", action='store_false') |
| parent_parser.add_argument('--no-use-encryption', |
| help='Do not us public key encryption for communication between the ' |
| 'randomizers and the analyzers via the shufflers. By default encryption ' |
| 'is used.', |
| dest="use_encryption", action='store_false') |
| |
| subparsers = parser.add_subparsers() |
| |
| sub_parser = subparsers.add_parser('build', parents=[parent_parser], |
| help='Builds the Cobalt prototype pipeline.') |
| sub_parser.set_defaults(func=_build) |
| |
| sub_parser = subparsers.add_parser('run', parents=[parent_parser], |
| help='Runs the synthetic data generator, the straight ' |
| 'counting pipeline and the Cobalt prototype pipeline.' |
| 'This is equivalent to: clean, gen, randomize, shuffle, analyze, ' |
| 'visualize.') |
| sub_parser.set_defaults(func=_run_all) |
| |
| sub_parser = subparsers.add_parser('test', parents=[parent_parser], |
| help='Runs the Cobalt tests.') |
| sub_parser.set_defaults(func=_test) |
| |
| sub_parser = subparsers.add_parser('clean', parents=[parent_parser], |
| help='Deletes the out directory.') |
| sub_parser.set_defaults(func=_clean_all) |
| |
| sub_parser = subparsers.add_parser('gen', parents=[parent_parser], |
| help='Generates fake input data and runs the straight counting pipeline.') |
| sub_parser.set_defaults(func=_generate) |
| |
| sub_parser = subparsers.add_parser('randomize', parents=[parent_parser], |
| help='Runs all the randomizers in Cobalt prototype pipeline.') |
| sub_parser.set_defaults(func=_randomize) |
| |
| sub_parser = subparsers.add_parser('shuffle', parents=[parent_parser], |
| help='Runs all the shufflers in Cobalt prototype pipeline.') |
| sub_parser.set_defaults(func=_shuffle) |
| |
| sub_parser = subparsers.add_parser('analyze', parents=[parent_parser], |
| help='Runs all the analyzers in Cobalt prototype pipeline.') |
| sub_parser.set_defaults(func=_analyze) |
| |
| sub_parser = subparsers.add_parser('visualize', parents=[parent_parser], |
| help='Generates the visualization data from Cobalt prototype pipeline.') |
| sub_parser.set_defaults(func=_visualize) |
| |
| args = parser.parse_args() |
| global _verbose_count |
| _verbose_count = args.verbose_count |
| global _use_public_key_encryption |
| _use_public_key_encryption = args.use_encryption |
| _initLogging(_verbose_count) |
| |
| return args.func() |
| |
| |
| if __name__ == '__main__': |
| sys.exit(main()) |