From c9ab66365194fdd441618c22662ea2b976f019c3 Mon Sep 17 00:00:00 2001 From: n1nj4sec Date: Fri, 25 Sep 2015 19:36:16 +0200 Subject: [PATCH] colored output now works on windows --- pupy/modules/ps.py | 4 ++-- pupy/pupylib/PupyCmd.py | 42 +++++++++++++++++++++++++++++++++++------ 2 files changed, 38 insertions(+), 8 deletions(-) diff --git a/pupy/modules/ps.py b/pupy/modules/ps.py index 2ef743d8..9dbc171b 100644 --- a/pupy/modules/ps.py +++ b/pupy/modules/ps.py @@ -2,9 +2,9 @@ from pupylib.PupyModule import * from pupylib.utils import obtain -__class_name__="MsgBoxPopup" +__class_name__="PsModule" -class MsgBoxPopup(PupyModule): +class PsModule(PupyModule): """ list processes """ def init_argparse(self): diff --git a/pupy/pupylib/PupyCmd.py b/pupy/pupylib/PupyCmd.py index ff8383a1..de8525f1 100644 --- a/pupy/pupylib/PupyCmd.py +++ b/pupy/pupylib/PupyCmd.py @@ -21,6 +21,7 @@ import re import os import os.path import traceback +import platform try: import ConfigParser as configparser except ImportError: @@ -57,6 +58,7 @@ BANNER=""" """%__version__ + def color_real(s, color, prompt=False, colors_enabled=True): """ color a string using ansi escape characters. set prompt to true to add marks for readline to see invisible portions of the prompt cf. http://stackoverflow.com/questions/9468435/look-how-to-fix-column-calculation-in-python-readline-if-use-color-prompt""" @@ -120,6 +122,23 @@ def obj2utf8(obj): obj=str(obj) return obj +class WindowsColoredStdout(object): + def __init__(self, write_color): + from ctypes import c_ulong, windll + STD_OUTPUT_HANDLE_ID = c_ulong(0xfffffff5) + windll.Kernel32.GetStdHandle.restype = c_ulong + self.std_output_hdl = windll.Kernel32.GetStdHandle(STD_OUTPUT_HANDLE_ID) + self.SetConsoleTextAttribute=windll.Kernel32.SetConsoleTextAttribute + self.write_color=write_color + def write(self, msg): + for attr, chunk in self.write_color(msg)[1]: + self.SetConsoleTextAttribute(self.std_output_hdl, attr.get_winattr()) + sys.stdout.write(chunk) + def flush(self): + sys.stdout.flush() + def read(self, *args, **kwargs): + sys.stdout.read(*args, **kwargs) + class PupyCmd(cmd.Cmd): def __init__(self, pupsrv, configFile="pupy.conf"): cmd.Cmd.__init__(self) @@ -133,6 +152,17 @@ class PupyCmd(cmd.Cmd): color = partial(color_real, colors_enabled=self.config.getboolean("cmdline","colors")) except Exception: color = color_real + + #wrap stdout to support ANSI coloring + if "windows" in platform.system().lower(): + if sys.stdout.isatty(): + try: + from pyreadline.console.ansi import write_color + self.stdout=WindowsColoredStdout(write_color) + except ImportError: + color = partial(color_real, colors_enabled=False) + self.display_warning("pyreadline is not installer. Output color disabled. Use \"pip install pyreadline\"") + self.intro = color(BANNER, 'green') self.prompt = color('>> ','blue', prompt=True) self.doc_header = 'Available commands :\n' @@ -335,18 +365,18 @@ class PupyCmd(cmd.Cmd): msg=str(msg) if msg: if modifier=="error": - sys.stdout.write(PupyCmd.format_error(msg)) + self.stdout.write(PupyCmd.format_error(msg)) elif modifier=="success": - sys.stdout.write(PupyCmd.format_success(msg)) + self.stdout.write(PupyCmd.format_success(msg)) elif modifier=="info": - sys.stdout.write(PupyCmd.format_info(msg)) + self.stdout.write(PupyCmd.format_info(msg)) elif modifier=="srvinfo": - sys.stdout.write(PupyCmd.format_srvinfo(msg)) + self.stdout.write(PupyCmd.format_srvinfo(msg)) #readline.redisplay() elif modifier=="warning": - sys.stdout.write(PupyCmd.format_warning(msg)) + self.stdout.write(PupyCmd.format_warning(msg)) else: - sys.stdout.write(PupyCmd.format_log(msg)) + self.stdout.write(PupyCmd.format_log(msg)) def display_srvinfo(self, msg): return self.display(msg, modifier="srvinfo")