Peter M. Groen / oletools

Commit 6ec331df7ab11f1df4ad232557752fb97fa7ab1e

Authored by Christian Herdtweck
1 parent 7a143086

msodde: handle case where sys.stdout cannot handle unicode 

This is the case e.g. if piping the output of msodde.py to file or other
program (e.g. grep) on a linux shell. In that case sys.stdout (i.e. print)
can only handle ascii, which is a problem when printing unicode dde links.
Inline Side-by-side
Showing 1 changed file with 50 additions and 0 deletions
oletools/msodde.py
  View file @6ec331d
... ... @@ -163,6 +163,56 @@ def get_logger(name, level=logging.CRITICAL+1):
163 163 log = get_logger('msodde')
164 164  
165 165  
  166 +# === UNICODE IN PY2 =========================================================
  167 +
  168 +def ensure_stdout_handles_unicode():
  169 + """ Ensure stdout can handle unicode by wrapping it if necessary
  170 +
  171 + Required e.g. if output of this script is piped or redirected in a linux
  172 + shell, since then sys.stdout.encoding is ascii and cannot handle
  173 + print(unicode). In that case we need to find some compatible encoding and
  174 + wrap sys.stdout into a encoder following (many thanks!)
  175 + https://stackoverflow.com/a/1819009 or https://stackoverflow.com/a/20447935
  176 +
  177 + Can be undone by setting sys.stdout = sys.__stdout__
  178 + """
  179 + import codecs
  180 + import locale
  181 +
  182 + # do not re-wrap
  183 + if isinstance(sys.stdout, codecs.StreamWriter):
  184 + return
  185 +
  186 + # try to find encoding for sys.stdout
  187 + encoding = None
  188 + try:
  189 + encoding = sys.stdout.encoding # variable encoding might not exist
  190 + except Exception:
  191 + pass
  192 +
  193 + if encoding not in (None, '', 'ascii'):
  194 + return # no need to wrap
  195 +
  196 + # try to find an encoding that can handle unicode
  197 + try:
  198 + encoding = locale.getpreferredencoding()
  199 + except Exception:
  200 + pass
  201 +
  202 + # fallback if still no encoding available
  203 + if encoding in (None, '', 'ascii'):
  204 + encoding = 'utf8'
  205 +
  206 + # logging is probably not initialized yet, but just in case
  207 + log.debug('wrapping sys.stdout with encoder using {0}'.format(encoding))
  208 +
  209 + wrapper = codecs.getwriter(encoding)
  210 + sys.stdout = wrapper(sys.stdout)
  211 +
  212 +
  213 +ensure_stdout_handles_unicode() # e.g. for print(text) in main()
  214 +
  215 +
166 216 # === ARGUMENT PARSING =======================================================
167 217  
168 218 class ArgParserWithBanner(argparse.ArgumentParser):
... ...