diff --git a/setup.py b/setup.py index 0d147b6..ea53c5e 100755 --- a/setup.py +++ b/setup.py @@ -35,7 +35,7 @@ include_package_data=True, scripts=['workbench/server/workbench_server', 'workbench_apps/workbench_cli/workbench'], tests_require=['tox'], - install_requires=['cython', 'colorama', 'elasticsearch', 'funcsigs', 'flask', 'filemagic', + install_requires=['cython', 'elasticsearch', 'funcsigs', 'flask', 'filemagic', 'ipython', 'lz4', 'mock', 'pandas', 'pefile', 'py2neo', 'pymongo', 'pytest', 'rekall', 'requests', 'ssdeep==2.9-0.3', 'urllib3', 'yara', 'zerorpc', 'cython'], diff --git a/workbench/server/plugin_manager.py b/workbench/server/plugin_manager.py index 60b04d2..9b0a117 100644 --- a/workbench/server/plugin_manager.py +++ b/workbench/server/plugin_manager.py @@ -8,8 +8,7 @@ from datetime import datetime import dir_watcher import inspect -import colorama -from colorama import Fore +from IPython.utils.coloransi import TermColors as color class PluginManager(object): """Plugin Manager for Workbench.""" @@ -83,8 +82,9 @@ def remove_plugin(self, f): """ if f.endswith('.py'): plugin_name = os.path.splitext(os.path.basename(f))[0] - print '- %s %sREMOVED' % (plugin_name, Fore.RED) - print '\t%sNote: still in memory, restart Workbench to remove...%s' % (Fore.YELLOW, Fore.RESET) + print '- %s %sREMOVED' % (plugin_name, color.Red) + print '\t%sNote: still in memory, restart Workbench to remove...%s' % \ + (color.Yellow, color.Normal) def add_plugin(self, f): """Adding and verifying plugin. @@ -101,7 +101,7 @@ def add_plugin(self, f): if plugin_name in sys.modules: try: handler = reload(sys.modules[plugin_name]) - print'\t- %s %sRELOAD%s' % (plugin_name, Fore.YELLOW, Fore.RESET) + print'\t- %s %sRELOAD%s' % (plugin_name, color.Yellow, color.Normal) except ImportError, error: print 'Failed to import plugin: %s (%s)' % (plugin_name, error) return @@ -115,7 +115,7 @@ def add_plugin(self, f): # Run the handler through plugin validation plugin = self.validate(handler) - print '\t- %s %sOK%s' % (plugin_name, Fore.GREEN, Fore.RESET) + print '\t- %s %sOK%s' % (plugin_name, color.Green, color.Normal) if plugin: # Okay must be successfully loaded so capture the plugin meta-data, diff --git a/workbench/server/workbench_server.py b/workbench/server/workbench_server.py index acc4bc5..37f91b4 100644 --- a/workbench/server/workbench_server.py +++ b/workbench/server/workbench_server.py @@ -18,9 +18,9 @@ import funcsigs import ConfigParser import magic -from colorama import Fore as F, Style import datetime import lz4 +from IPython.utils.coloransi import TermColors as color try: from cStringIO import StringIO except ImportError: @@ -620,27 +620,27 @@ def help(self, topic=None): # message that has both the md5 of what they were looking for and # a nice informative message that explains what might have happened sample_md5 = e.args[0] - return '%s%s\n\t%s%s%s' % (F.YELLOW, sample_md5, F.GREEN, e.message(), F.RESET) + return '%s%s\n\t%s%s%s' % (color.Yellow, sample_md5, color.Green, e.message(), color.Normal) # Fixme: These are internal methods that basically just provide help text def _help_workbench(self): """ Help on Workbench """ - help = '%sWelcome to Workbench Help:%s' % (F.YELLOW, F.RESET) - help += '\n\t%s- workbench.help(\'basic\') %s for getting started help' % (F.GREEN, F.BLUE) - help += '\n\t%s- workbench.help(\'workers\') %s for help on available workers' % (F.GREEN, F.BLUE) - help += '\n\t%s- workbench.help(\'commands\') %s for help on workbench commands' % (F.GREEN, F.BLUE) - help += '\n\t%s- workbench.help(topic) %s where topic can be a help, command or worker' % (F.GREEN, F.BLUE) - help += '\n\n%sSee http://github.com/SuperCowPowers/workbench for more information\n%s' % (F.YELLOW, F.RESET) + help = '%sWelcome to Workbench Help:%s' % (color.Yellow, color.Normal) + help += '\n\t%s- workbench.help(\'basic\') %s for getting started help' % (color.Green, color.Blue) + help += '\n\t%s- workbench.help(\'workers\') %s for help on available workers' % (color.Green, color.Blue) + help += '\n\t%s- workbench.help(\'commands\') %s for help on workbench commands' % (color.Green, color.Blue) + help += '\n\t%s- workbench.help(topic) %s where topic can be a help, command or worker' % (color.Green, color.Blue) + help += '\n\n%sSee http://github.com/SuperCowPowers/workbench for more information\n%s' % (color.Yellow, color.Normal) return help def _help_basic(self): """ Help for Workbench Basics """ - help = '%sWorkbench: Getting started...' % (F.YELLOW) - help += '\n%sStore a sample into Workbench:' % (F.GREEN) - help += '\n\t%s$ workbench.store_sample(raw_bytes, filename, type_tag)' % (F.BLUE) - help += '\n\n%sNotice store_sample returns an md5 of the sample...'% (F.YELLOW) - help += '\n%sRun workers on the sample (view, meta, whatever...):' % (F.GREEN) - help += '\n\t%s$ workbench.work_request(\'view\', md5)%s' % (F.BLUE, F.RESET) + help = '%sWorkbench: Getting started...' % (color.Yellow) + help += '\n%sStore a sample into Workbench:' % (color.Green) + help += '\n\t%s$ workbench.store_sample(raw_bytes, filename, type_tag)' % (color.Blue) + help += '\n\n%sNotice store_sample returns an md5 of the sample...'% (color.Yellow) + help += '\n%sRun workers on the sample (view, meta, whatever...):' % (color.Green) + help += '\n\t%s$ workbench.work_request(\'view\', md5)%s' % (color.Blue, color.Normal) return help def _help_commands(self): diff --git a/workbench/workers/help_formatter.py b/workbench/workers/help_formatter.py index 8712cd1..d29fa3d 100644 --- a/workbench/workers/help_formatter.py +++ b/workbench/workers/help_formatter.py @@ -1,7 +1,7 @@ ''' HelpFormatter worker ''' -from colorama import Fore, Style +from IPython.utils.coloransi import TermColors as color class HelpFormatter(object): ''' This worker does CLI formatting and coloring for any help object ''' @@ -14,24 +14,23 @@ def execute(self, input_data): # Standard help text if type_tag == 'help': - output = '%s%s%s%s%s' % (Style.BRIGHT, Fore.BLUE, input_data['help'], Fore.RESET, Style.RESET_ALL) + output = '%s%s%s' % (color.Blue, input_data['help'], color.Normal) # Worker elif type_tag == 'worker': - output = '%s%s%s%s' % (Style.BRIGHT, Fore.YELLOW, input_data['name'], Style.RESET_ALL) - output += '\n %sInput: %s%s%s' % (Fore.BLUE, Fore.GREEN, input_data['dependencies'], Fore.RESET) - output += '\n %s%s' % (Fore.GREEN, input_data['docstring']) + output = '%s%s' % (color.Yellow, input_data['name']) + output += '\n %sInput: %s%s%s' % (color.Blue, color.Green, input_data['dependencies'], color.Normal) + output += '\n %s%s' % (color.Green, input_data['docstring']) # Command elif type_tag == 'command': - output = '%s%s%s%s%s %s' % (Style.BRIGHT, Fore.YELLOW, input_data['command'], - Style.RESET_ALL, Fore.BLUE, input_data['sig']) - output += '\n %s%s%s' % (Fore.GREEN, input_data['docstring'], Fore.RESET) + output = '%s%s%s %s' % (color.Yellow, input_data['command'], color.Blue, input_data['sig']) + output += '\n %s%s%s' % (color.Green, input_data['docstring'], color.Normal) # WTF: Alert on unknown type_tag and return a string of the input_data else: print 'Alert: help_formatter worker received malformed object: %s' % str(input_data) - output = '\n%s%s%s' % (Fore.RED, str(input_data), Fore.RESET) + output = '\n%s%s%s' % (color.Red, str(input_data), color.Normal) # Return the formatted and colored help return {'help': output} diff --git a/workbench_apps/setup.py b/workbench_apps/setup.py index b475a14..7069274 100755 --- a/workbench_apps/setup.py +++ b/workbench_apps/setup.py @@ -30,7 +30,7 @@ package_dir={'workbench': 'workbench_cli'}, include_package_data=True, scripts=['workbench_cli/workbench'], - install_requires=['colorama', 'funcsigs', 'ipython', 'lz4', + install_requires=['funcsigs', 'ipython', 'lz4', 'pandas', 'pytest', 'zerorpc'], license='MIT', zip_safe=False, diff --git a/workbench_apps/workbench_cli/help_content.py b/workbench_apps/workbench_cli/help_content.py index 6371ec0..a40a431 100644 --- a/workbench_apps/workbench_cli/help_content.py +++ b/workbench_apps/workbench_cli/help_content.py @@ -1,7 +1,7 @@ """Workbench Interactive Shell Help Content""" import inspect -from colorama import Fore as F +from IPython.utils.coloransi import TermColors as color class WorkbenchShellHelp(object): """Workbench CLI Help""" @@ -11,87 +11,87 @@ def __init__(self): def help_cli(self): """ Help on Workbench CLI """ - help = '%sWelcome to Workbench CLI Help:%s' % (F.YELLOW, F.RESET) - help += '\n\t%s> help cli_basic %s for getting started help' % (F.GREEN, F.BLUE) - help += '\n\t%s> help workers %s for help on available workers' % (F.GREEN, F.BLUE) - help += '\n\t%s> help search %s for help on searching samples' % (F.GREEN, F.BLUE) - help += '\n\t%s> help dataframe %s for help on making dataframes' % (F.GREEN, F.BLUE) - help += '\n\t%s> help commands %s for help on workbench commands' % (F.GREEN, F.BLUE) - help += '\n\t%s> help topic %s where topic can be a help, command or worker' % (F.GREEN, F.BLUE) - help += '\n\n%sNote: cli commands are transformed into python calls' % (F.YELLOW) - help += '\n\t%s> help cli_basic --> help("cli_basic")%s' % (F.GREEN, F.RESET) + help = '%sWelcome to Workbench CLI Help:%s' % (color.Yellow, color.Normal) + help += '\n\t%s> help cli_basic %s for getting started help' % (color.Green, color.Blue) + help += '\n\t%s> help workers %s for help on available workers' % (color.Green, color.Blue) + help += '\n\t%s> help search %s for help on searching samples' % (color.Green, color.Blue) + help += '\n\t%s> help dataframe %s for help on making dataframes' % (color.Green, color.Blue) + help += '\n\t%s> help commands %s for help on workbench commands' % (color.Green, color.Blue) + help += '\n\t%s> help topic %s where topic can be a help, command or worker' % (color.Green, color.Blue) + help += '\n\n%sNote: cli commands are transformed into python calls' % (color.Yellow) + help += '\n\t%s> help cli_basic --> help("cli_basic")%s' % (color.Green, color.Normal) return help def help_cli_basic(self): """ Help for Workbench CLI Basics """ - help = '%sWorkbench: Getting started...' % (F.YELLOW) - help += '\n%sLoad in a sample:' % (F.GREEN) - help += '\n\t%s> load_sample /path/to/file' % (F.BLUE) - help += '\n\n%sNotice the prompt now shows the md5 of the sample...'% (F.YELLOW) - help += '\n%sRun workers on the sample:' % (F.GREEN) - help += '\n\t%s> view' % (F.BLUE) - help += '\n%sType the \'help workers\' or the first part of the worker ...' % (F.GREEN) - help += '\n\t%s> help workers (lists all possible workers)' % (F.BLUE) - help += '\n\t%s> pe_ (will give you pe_classifier, pe_deep_sim, pe_features, pe_indicators, pe_peid)%s' % (F.BLUE, F.RESET) + help = '%sWorkbench: Getting started...' % (color.Yellow) + help += '\n%sLoad in a sample:' % (color.Green) + help += '\n\t%s> load_sample /path/to/file' % (color.Blue) + help += '\n\n%sNotice the prompt now shows the md5 of the sample...'% (color.Yellow) + help += '\n%sRun workers on the sample:' % (color.Green) + help += '\n\t%s> view' % (color.Blue) + help += '\n%sType the \'help workers\' or the first part of the worker ...' % (color.Green) + help += '\n\t%s> help workers (lists all possible workers)' % (color.Blue) + help += '\n\t%s> pe_ (will give you pe_classifier, pe_deep_sim, pe_features, pe_indicators, pe_peid)%s' % (color.Blue, color.Normal) return help def help_cli_search(self): """ Help for Workbench CLI Search """ - help = '%sSearch: %s returns sample_sets, a sample_set is a set/list of md5s.' % (F.YELLOW, F.GREEN) - help += '\n\n\t%sSearch for all samples in the database that are known bad pe files,' % (F.GREEN) - help += '\n\t%sthis command returns the sample_set containing the matching items'% (F.GREEN) - help += '\n\t%s> my_bad_exes = search([\'bad\', \'exe\'])' % (F.BLUE) - help += '\n\n\t%sRun workers on this sample_set:' % (F.GREEN) - help += '\n\t%s> pe_outputs = pe_features(my_bad_exes) %s' % (F.BLUE, F.RESET) - help += '\n\n\t%sLoop on the generator (or make a DataFrame see >help dataframe)' % (F.GREEN) - help += '\n\t%s> for output in pe_outputs: %s' % (F.BLUE, F.RESET) - help += '\n\t\t%s print output %s' % (F.BLUE, F.RESET) + help = '%sSearch: %s returns sample_sets, a sample_set is a set/list of md5s.' % (color.Yellow, color.Green) + help += '\n\n\t%sSearch for all samples in the database that are known bad pe files,' % (color.Green) + help += '\n\t%sthis command returns the sample_set containing the matching items'% (color.Green) + help += '\n\t%s> my_bad_exes = search([\'bad\', \'exe\'])' % (color.Blue) + help += '\n\n\t%sRun workers on this sample_set:' % (color.Green) + help += '\n\t%s> pe_outputs = pe_features(my_bad_exes) %s' % (color.Blue, color.Normal) + help += '\n\n\t%sLoop on the generator (or make a DataFrame see >help dataframe)' % (color.Green) + help += '\n\t%s> for output in pe_outputs: %s' % (color.Blue, color.Normal) + help += '\n\t\t%s print output %s' % (color.Blue, color.Normal) return help def help_dataframe(self): """ Help for making a DataFrame with Workbench CLI """ - help = '%sMaking a DataFrame: %s how to make a dataframe from raw data (pcap, memory, pe files)' % (F.YELLOW, F.GREEN) - help += '\n\t%sNote: for memory_image and pe_files see > help dataframe_memory or dataframe_pe' % (F.GREEN) - help += '\n\n%sPCAP Example:' % (F.GREEN) - help += '\n\t%s> load_sample /path/to/pcap/gold_xxx.pcap [\'bad\', \'threatglass\']' % (F.BLUE) - help += '\n\t%s> view # view is your friend use it often' % (F.BLUE) - help += '\n\n%sGrab the http_log from the pcap (also play around with other logs):' % (F.GREEN) - help += '\n\t%s> http_log_md5 = view()[\'view\'][\'bro_logs\'][\'http_log\']' % (F.BLUE) - help += '\n\t%s> http_log_md5 (returns the md5 of the http_log)' % (F.BLUE) - help += '\n\n%sStream back the ^contents^ of the http_log:' % (F.GREEN) - help += '\n\t%s> http_log = stream_sample(http_log_md5)' % (F.BLUE) - help += '\n\n%sPut the http_log into a dataframe:' % (F.GREEN) - help += '\n\t%s> http_df = pd.DataFrame(http_log)' % (F.BLUE) - help += '\n\t%s> http_df.head()' % (F.BLUE) - help += '\n\t%s> http_df.groupby([\'host\',\'id.resp_h\',\'resp_mime_types\'])[[\'response_body_len\']].sum()' % (F.BLUE) - help += '\n\t%s> http_df.describe() %s' % (F.BLUE, F.RESET) + help = '%sMaking a DataFrame: %s how to make a dataframe from raw data (pcap, memory, pe files)' % (color.Yellow, color.Green) + help += '\n\t%sNote: for memory_image and pe_files see > help dataframe_memory or dataframe_pe' % (color.Green) + help += '\n\n%sPCAP Example:' % (color.Green) + help += '\n\t%s> load_sample /path/to/pcap/gold_xxx.pcap [\'bad\', \'threatglass\']' % (color.Blue) + help += '\n\t%s> view # view is your friend use it often' % (color.Blue) + help += '\n\n%sGrab the http_log from the pcap (also play around with other logs):' % (color.Green) + help += '\n\t%s> http_log_md5 = view()[\'view\'][\'bro_logs\'][\'http_log\']' % (color.Blue) + help += '\n\t%s> http_log_md5 (returns the md5 of the http_log)' % (color.Blue) + help += '\n\n%sStream back the ^contents^ of the http_log:' % (color.Green) + help += '\n\t%s> http_log = stream_sample(http_log_md5)' % (color.Blue) + help += '\n\n%sPut the http_log into a dataframe:' % (color.Green) + help += '\n\t%s> http_df = pd.DataFrame(http_log)' % (color.Blue) + help += '\n\t%s> http_df.head()' % (color.Blue) + help += '\n\t%s> http_df.groupby([\'host\',\'id.resp_h\',\'resp_mime_types\'])[[\'response_body_len\']].sum()' % (color.Blue) + help += '\n\t%s> http_df.describe() %s' % (color.Blue, color.Normal) return help def help_dataframe_memory(self): """ Help for making a DataFrame with Workbench CLI """ - help = '%sMaking a DataFrame: %s how to make a dataframe from memory_forensics sample' % (F.YELLOW, F.GREEN) - help += '\n\n%sMemory Images Example:' % (F.GREEN) - help += '\n\t%s> load_sample /path/to/pcap/exemplar4.vmem [\'bad\', \'aptz13\']' % (F.BLUE) - help += '\n\t%s> view # view is your friend use it often' % (F.BLUE) - help += '\n\t%s> <<< TODO :) >>> %s' % (F.BLUE, F.RESET) + help = '%sMaking a DataFrame: %s how to make a dataframe from memory_forensics sample' % (color.Yellow, color.Green) + help += '\n\n%sMemory Images Example:' % (color.Green) + help += '\n\t%s> load_sample /path/to/pcap/exemplar4.vmem [\'bad\', \'aptz13\']' % (color.Blue) + help += '\n\t%s> view # view is your friend use it often' % (color.Blue) + help += '\n\t%s> <<< TODO :) >>> %s' % (color.Blue, color.Normal) return help def help_dataframe_pe(self): """ Help for making a DataFrame with Workbench CLI """ - help = '%sMaking a DataFrame: %s how to make a dataframe from pe files' % (F.YELLOW, F.GREEN) - help += '\n\n%sPE Files Example (loading a directory):' % (F.GREEN) - help += '\n\t%s> load_sample /path/to/pe/bad [\'bad\', \'case_69\']' % (F.BLUE) - help += '\n\n\t%sSearch for all samples in the database that are pe files,' % (F.GREEN) - help += '\n\t%sthis command returns the sample_set containing the matching items'% (F.GREEN) - help += '\n\t%s> my_exes = search([\'exe\'])' % (F.BLUE) - help += '\n\n\t%sRun workers on this sample_set:' % (F.GREEN) - help += '\n\t%s> pe_outputs = set_work_request(\'pe_features\', my_exes, [\'md5\', \'dense_features.*\', \'tags\'])' % (F.BLUE) - help += '\n\n\t%sMake a DataFrame:' % (F.GREEN) - help += '\n\t%s> pe_df = pd.DataFrame(pe_outputs) %s' % (F.BLUE, F.RESET) - help += '\n\t%s> pe_df.head() %s' % (F.BLUE, F.RESET) - help += '\n\t%s> pe_df = flatten_tags(pe_df) %s' % (F.BLUE, F.RESET) - help += '\n\t%s> pe_df.hist(\'check_sum\',\'tags\') %s' % (F.BLUE, F.RESET) - help += '\n\t%s> pe_df.bloxplot(\'check_sum\',\'tags\') %s' % (F.BLUE, F.RESET) + help = '%sMaking a DataFrame: %s how to make a dataframe from pe files' % (color.Yellow, color.Green) + help += '\n\n%sPE Files Example (loading a directory):' % (color.Green) + help += '\n\t%s> load_sample /path/to/pe/bad [\'bad\', \'case_69\']' % (color.Blue) + help += '\n\n\t%sSearch for all samples in the database that are pe files,' % (color.Green) + help += '\n\t%sthis command returns the sample_set containing the matching items'% (color.Green) + help += '\n\t%s> my_exes = search([\'exe\'])' % (color.Blue) + help += '\n\n\t%sRun workers on this sample_set:' % (color.Green) + help += '\n\t%s> pe_outputs = set_work_request(\'pe_features\', my_exes, [\'md5\', \'dense_features.*\', \'tags\'])' % (color.Blue) + help += '\n\n\t%sMake a DataFrame:' % (color.Green) + help += '\n\t%s> pe_df = pd.DataFrame(pe_outputs) %s' % (color.Blue, color.Normal) + help += '\n\t%s> pe_df.head() %s' % (color.Blue, color.Normal) + help += '\n\t%s> pe_df = flatten_tags(pe_df) %s' % (color.Blue, color.Normal) + help += '\n\t%s> pe_df.hist(\'check_sum\',\'tags\') %s' % (color.Blue, color.Normal) + help += '\n\t%s> pe_df.bloxplot(\'check_sum\',\'tags\') %s' % (color.Blue, color.Normal) return help ################## @@ -108,7 +108,7 @@ def test(): # Now execute all the help methods for name, method in help._all_help_methods().iteritems(): - print '\n%s%s%s' % (F.RED, name, F.RESET) + print '\n%s%s%s' % (color.Red, name, color.Normal) print '%s' % method() if __name__ == '__main__': diff --git a/workbench_apps/workbench_cli/workbench_shell.py b/workbench_apps/workbench_cli/workbench_shell.py index 37c18ef..c3ddf21 100644 --- a/workbench_apps/workbench_cli/workbench_shell.py +++ b/workbench_apps/workbench_cli/workbench_shell.py @@ -8,21 +8,21 @@ import inspect import funcsigs import operator -from colorama import Fore as F import pprint +from IPython.utils.coloransi import TermColors as color try: import pandas as pd except ImportError: - print '\n%sNotice: pandas not found...' % F.YELLOW - print '\t%sWe recommend installing pandas: %s$ pip install pandas%s' % (F.BLUE, F.RED, F.RESET) + print '\n%sNotice: pandas not found...' % color.Yellow + print '\t%sWe recommend installing pandas: %s$ pip install pandas%s' % (color.Blue, color.Red, color.Normal) try: import matplotlib.pyplot as plt plt.ion() except ImportError: - print '\n%sNotice: matplotlib not found...' % F.YELLOW - print '\t%sWe recommend installing matplotlib: %s$ pip install matplotlib%s' % (F.BLUE, F.RED, F.RESET) + print '\n%sNotice: matplotlib not found...' % color.Yellow + print '\t%sWe recommend installing matplotlib: %s$ pip install matplotlib%s' % (color.Blue, color.Red, color.Normal) try: from . import client_helper from . import help_content @@ -96,7 +96,7 @@ def load_sample(self, file_path, tags=None): # Recommend a tag if not tags: print '\n%sRecommended: Add a list of tags when you load samples. \ - \n\t%sExamples: [\'bad\'], [\'good\'], [\'bad\',\'aptz13\']%s' % (F.YELLOW, F.GREEN, F.RESET) + \n\t%sExamples: [\'bad\'], [\'good\'], [\'bad\',\'aptz13\']%s' % (color.Yellow, color.Green, color.Normal) return # Do they want everything under a directory? @@ -111,12 +111,12 @@ def load_sample(self, file_path, tags=None): raw_bytes = my_file.read() md5 = hashlib.md5(raw_bytes).hexdigest() if not self.workbench.has_sample(md5): - print '%sStreaming Sample...%s' % (F.MAGENTA, F.RESET) + print '%sStreaming Sample...%s' % (color.Purple, color.Normal) basename = os.path.basename(path) md5 = self.streamer.stream_to_workbench(raw_bytes, basename, 'unknown', tags) print '\n%s %s%s %sLocked and Loaded...%s\n' % \ - (self.beer, F.MAGENTA, md5[:6], F.YELLOW, F.RESET) + (self.beer, color.Purple, md5[:6], color.Yellow, color.Normal) # Add tags to the sample self.workbench.add_tags(md5, tags) @@ -147,14 +147,14 @@ def tag_info(self): corr = tag_df.corr().fillna(1) corr_dict = corr.to_dict() - print '\n%sSamples in Database%s' % (F.MAGENTA, F.RESET) + print '\n%sSamples in Database%s' % (color.Purple, color.Normal) for tag, count in tag_freq.iteritems(): - print ' %s%s: %s%s%s (' % (F.GREEN, tag, F.BLUE, count, F.RESET), + print ' %s%s: %s%s%s (' % (color.Green, tag, color.Blue, count, color.Normal), tag_corrs = sorted(corr_dict[tag].iteritems(), key=operator.itemgetter(1), reverse=True) for corr_tag, value in tag_corrs[:5]: if corr_tag != tag: - print '%s%s:%s%.1f' % (F.GREEN, corr_tag, F.BLUE, value), - print '%s)' % F.RESET + print '%s%s:%s%.1f' % (color.Green, corr_tag, color.Blue, value), + print '%s)' % color.Normal def pull_df(self, md5): """Wrapper for the Workbench get_dataframe method @@ -193,7 +193,7 @@ def versions(self): Returns: The running versions of both the CLI and the Workbench Server """ - print '%s<<< Workbench CLI Version %s >>>%s' % (F.BLUE, self.version, F.RESET) + print '%s<<< Workbench CLI Version %s >>>%s' % (color.Blue, self.version, color.Normal) print self.workbench.help('version') def run(self): @@ -215,9 +215,7 @@ def run(self): cfg.InteractiveShellEmbed.autoindent = True cfg.InteractiveShellEmbed.deep_reload = True cfg.PromptManager.in_template = ( - r'{color.Purple}' - r'{short_md5}' - r'{color.Blue} Workbench{color.Green}[\#]> ') + r'{color.Purple}{short_md5}{color.Blue} Workbench{color.Green}[\#]> ') # cfg.PromptManager.out_template = '' # Create the IPython shell @@ -247,7 +245,7 @@ def _connect(self, server_info): del _tmp_connect except zerorpc.exceptions.LostRemote: print '%sError: Could not connect to Workbench Server at %s:%s%s' % \ - (F.RED, server_info['server'], server_info['port'], F.RESET) + (color.Red, server_info['server'], server_info['port'], color.Normal) sys.exit(1) # Okay do the real connection @@ -255,7 +253,7 @@ def _connect(self, server_info): self.workbench.close() self.workbench = zerorpc.Client(timeout=300, heartbeat=60) self.workbench.connect('tcp://'+server_info['server']+':'+server_info['port']) - print '\n%s<<< Connected: %s:%s >>>%s' % (F.GREEN, server_info['server'], server_info['port'], F.RESET) + print '\n%s<<< Connected: %s:%s >>>%s' % (color.Green, server_info['server'], server_info['port'], color.Normal) def _progress_print(self, sent, total): """Progress print show the progress of the current upload with a neat progress bar @@ -263,7 +261,7 @@ def _progress_print(self, sent, total): """ percent = min(int(sent*100.0/total), 100) sys.stdout.write('\r{0}[{1}{2}] {3}{4}%{5}'. - format(F.GREEN, '#'*(percent/2), ' '*(50-percent/2), F.YELLOW, percent, F.RESET)) + format(color.Green, '#'*(percent/2), ' '*(50-percent/2), color.Yellow, percent, color.Normal)) sys.stdout.flush() def _work_request(self, worker, md5=None): @@ -287,7 +285,7 @@ def _work_request(self, worker, md5=None): def _data_not_found(self, e): """Message when you get a DataNotFound exception from the server""" - return '%s%s%s' % (F.RED, e.msg, F.RESET) + return '%s%s%s' % (color.Red, e.msg, color.Normal) def _generate_command_dict(self): """Create a customized namespace for Workbench with a bunch of shortcuts