Peter M. Groen / oletools

Commit 77cb3e3bfbe5a219775a467f3ec26a1e5c5ddc8c

Authored by Christian Herdtweck
1 parent bda20c7b

olevba[3]: move arg-parsing to own function with optional custom args 

This simplifies unit test creation a lot
Inline Side-by-side
Showing 2 changed files with 42 additions and 12 deletions
oletools/olevba.py
  View file @77cb3e3
@@ -3269,10 +3269,9 @@ class VBA_Parser_CLI(VBA_Parser): @@ -3269,10 +3269,9 @@ class VBA_Parser_CLI(VBA_Parser):
3269 3269
3270 #=== MAIN ===================================================================== 3270 #=== MAIN =====================================================================
3271 3271
3272 -def main():  
3273 - """  
3274 - Main function, called when olevba is run from the command line  
3275 - """ 3272 +def parse_args(cmd_line_args=None):
  3273 + """ parse command line arguments (given ones or per default sys.argv) """
  3274 +
3276 DEFAULT_LOG_LEVEL = "warning" # Default log level 3275 DEFAULT_LOG_LEVEL = "warning" # Default log level
3277 LOG_LEVELS = { 3276 LOG_LEVELS = {
3278 'debug': logging.DEBUG, 3277 'debug': logging.DEBUG,
@@ -3324,7 +3323,7 @@ def main(): @@ -3324,7 +3323,7 @@ def main():
3324 parser.add_option('--relaxed', dest="relaxed", action="store_true", default=False, 3323 parser.add_option('--relaxed', dest="relaxed", action="store_true", default=False,
3325 help="Do not raise errors if opening of substream fails") 3324 help="Do not raise errors if opening of substream fails")
3326 3325
3327 - (options, args) = parser.parse_args() 3326 + (options, args) = parser.parse_args(cmd_line_args)
3328 3327
3329 # Print help if no arguments are passed 3328 # Print help if no arguments are passed
3330 if len(args) == 0: 3329 if len(args) == 0:
@@ -3333,6 +3332,22 @@ def main(): @@ -3333,6 +3332,22 @@ def main():
3333 parser.print_help() 3332 parser.print_help()
3334 sys.exit(RETURN_WRONG_ARGS) 3333 sys.exit(RETURN_WRONG_ARGS)
3335 3334
  3335 + options.loglevel = LOG_LEVELS[options.loglevel]
  3336 +
  3337 + return options, args
  3338 +
  3339 +
  3340 +def main(cmd_line_args=None):
  3341 + """
  3342 + Main function, called when olevba is run from the command line
  3343 +
  3344 + Optional argument: command line arguments to be forwarded to ArgumentParser
  3345 + in process_args. Per default (cmd_line_args=None), sys.argv is used. Option
  3346 + mainly added for unit-testing
  3347 + """
  3348 +
  3349 + options, args = parse_args(cmd_line_args)
  3350 +
3336 # provide info about tool and its version 3351 # provide info about tool and its version
3337 if options.output_mode == 'json': 3352 if options.output_mode == 'json':
3338 # prints opening [ 3353 # prints opening [
@@ -3342,7 +3357,7 @@ def main(): @@ -3342,7 +3357,7 @@ def main():
3342 else: 3357 else:
3343 print('olevba %s - http://decalage.info/python/oletools' % __version__) 3358 print('olevba %s - http://decalage.info/python/oletools' % __version__)
3344 3359
3345 - logging.basicConfig(level=LOG_LEVELS[options.loglevel], format='%(levelname)-8s %(message)s') 3360 + logging.basicConfig(level=options.loglevel, format='%(levelname)-8s %(message)s')
3346 # enable logging in the modules: 3361 # enable logging in the modules:
3347 enable_logging() 3362 enable_logging()
3348 3363
oletools/olevba3.py
  View file @77cb3e3
@@ -3232,10 +3232,9 @@ class VBA_Parser_CLI(VBA_Parser): @@ -3232,10 +3232,9 @@ class VBA_Parser_CLI(VBA_Parser):
3232 3232
3233 #=== MAIN ===================================================================== 3233 #=== MAIN =====================================================================
3234 3234
3235 -def main():  
3236 - """  
3237 - Main function, called when olevba is run from the command line  
3238 - """ 3235 +def parse_args(cmd_line_args=None):
  3236 + """ parse command line arguments (given ones or per default sys.argv) """
  3237 +
3239 DEFAULT_LOG_LEVEL = "warning" # Default log level 3238 DEFAULT_LOG_LEVEL = "warning" # Default log level
3240 LOG_LEVELS = { 3239 LOG_LEVELS = {
3241 'debug': logging.DEBUG, 3240 'debug': logging.DEBUG,
@@ -3287,7 +3286,7 @@ def main(): @@ -3287,7 +3286,7 @@ def main():
3287 parser.add_option('--relaxed', dest="relaxed", action="store_true", default=False, 3286 parser.add_option('--relaxed', dest="relaxed", action="store_true", default=False,
3288 help="Do not raise errors if opening of substream fails") 3287 help="Do not raise errors if opening of substream fails")
3289 3288
3290 - (options, args) = parser.parse_args() 3289 + (options, args) = parser.parse_args(cmd_line_args)
3291 3290
3292 # Print help if no arguments are passed 3291 # Print help if no arguments are passed
3293 if len(args) == 0: 3292 if len(args) == 0:
@@ -3295,6 +3294,22 @@ def main(): @@ -3295,6 +3294,22 @@ def main():
3295 parser.print_help() 3294 parser.print_help()
3296 sys.exit(RETURN_WRONG_ARGS) 3295 sys.exit(RETURN_WRONG_ARGS)
3297 3296
  3297 + options.loglevel = LOG_LEVELS[options.loglevel]
  3298 +
  3299 + return options, args
  3300 +
  3301 +
  3302 +def main(cmd_line_args=None):
  3303 + """
  3304 + Main function, called when olevba is run from the command line
  3305 +
  3306 + Optional argument: command line arguments to be forwarded to ArgumentParser
  3307 + in process_args. Per default (cmd_line_args=None), sys.argv is used. Option
  3308 + mainly added for unit-testing
  3309 + """
  3310 +
  3311 + options, args = parse_args(cmd_line_args)
  3312 +
3298 # provide info about tool and its version 3313 # provide info about tool and its version
3299 if options.output_mode == 'json': 3314 if options.output_mode == 'json':
3300 # prints opening [ 3315 # prints opening [
@@ -3304,7 +3319,7 @@ def main(): @@ -3304,7 +3319,7 @@ def main():
3304 else: 3319 else:
3305 print('olevba %s - http://decalage.info/python/oletools' % __version__) 3320 print('olevba %s - http://decalage.info/python/oletools' % __version__)
3306 3321
3307 - logging.basicConfig(level=LOG_LEVELS[options.loglevel], format='%(levelname)-8s %(message)s') 3322 + logging.basicConfig(level=options.loglevel, format='%(levelname)-8s %(message)s')
3308 # enable logging in the modules: 3323 # enable logging in the modules:
3309 log.setLevel(logging.NOTSET) 3324 log.setLevel(logging.NOTSET)
3310 3325