Peter M. Groen / oletools

Commit ecd49a760d0e180160148bc4f8fb2b5c4759ed54

Authored by Christian Herdtweck
1 parent 8f9df350

Add option for custom command line args in msodde.py 

Custom command line arguments greatly simplifies unit test creation for
the main method.

Also, avoid calling sys.exit in functions (that is good practice if I
recall corretly).
Inline Side-by-side
Showing 1 changed file with 14 additions and 6 deletions
oletools/msodde.py
  View file @ecd49a7
... ... @@ -121,14 +121,14 @@ def existing_file(filename):
121 121 return filename
122 122  
123 123  
124   -def process_args():
  124 +def process_args(cmd_line_args=None):
125 125 parser = ArgParserWithBanner(description='A python tool to detect and extract DDE links in MS Office files')
126 126 parser.add_argument("filepath", help="path of the file to be analyzed",
127 127 type=existing_file, metavar='FILE')
128 128 parser.add_argument("--json", '-j', action='store_true',
129 129 help="Output in json format")
130 130  
131   - return parser.parse_args()
  131 + return parser.parse_args(cmd_line_args)
132 132  
133 133  
134 134 # === FUNCTIONS ==============================================================
... ... @@ -312,8 +312,14 @@ def process_file(filepath):
312 312  
313 313 #=== MAIN =================================================================
314 314  
315   -def main():
316   - args = process_args()
  315 +def main(cmd_line_args=None):
  316 + """ Main function, called if this file is called as a script
  317 +
  318 + Optional argument: command line arguments to be forwarded to ArgumentParser
  319 + in process_args. Per default (cmd_line_args=None), sys.argv is used. Option
  320 + mainly added for unit-testing
  321 + """
  322 + args = process_args(cmd_line_args)
317 323  
318 324 if args.json:
319 325 jout = []
... ... @@ -342,11 +348,13 @@ def main():
342 348 jout.append(dict(type='dde-link', link=line.strip()))
343 349 json.dump(jout, sys.stdout, check_circular=False, indent=4)
344 350 print() # add a newline after closing "]"
345   - sys.exit(return_code) # required if we catch an exception in json-mode
  351 + return return_code # required if we catch an exception in json-mode
346 352 else:
347 353 print ('DDE Links:')
348 354 print(text)
349 355  
  356 + return return_code
  357 +
350 358  
351 359 if __name__ == '__main__':
352   - main()
  360 + sys.exit(main())
... ...