From 86a14b53fb0b482692835a8f89b3a12017de0fdc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marek=20Marczykowski-G=C3=B3recki?=
 <marmarek@invisiblethingslab.com>
Date: Fri, 29 Jul 2016 11:20:21 +0200
Subject: [PATCH] qvm-run: color untrusted stderr even when stdout is
 redirected

When stdout is redirected to some file or command two things will
happen:
 - qvm-run will not automatically color the output as stdout is not a
 TTY
 - even when coloring is forced, it will not work, as the control
 sequence (on stdout) will be redirected anyway

Fix this by handling stdout and stderr independently and output color
switching sequence to each of them.

Fixes QubesOS/qubes-issues#2190
---
 qvm-tools/qvm-run | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/qvm-tools/qvm-run b/qvm-tools/qvm-run
index 83223fe2..e255fc19 100755
--- a/qvm-tools/qvm-run
+++ b/qvm-tools/qvm-run
@@ -47,7 +47,9 @@ def vm_run_cmd(vm, cmd, options):
     if options.verbose:
         print >> sys.stderr, "Running command on VM: '{0}'...".format(vm.name)
     if options.passio and options.color_output is not None:
-        print "\033[0;%dm" % options.color_output,
+        sys.stdout.write("\033[0;{}m".format(options.color_output))
+    if options.passio and options.color_stderr is not None:
+        sys.stderr.write("\033[0;{}m".format(options.color_stderr))
 
     try:
         def tray_notify_generic(level, str):
@@ -65,6 +67,8 @@ def vm_run_cmd(vm, cmd, options):
     except QubesException as err:
         if options.passio and options.color_output is not None:
             sys.stdout.write("\033[0m")
+        if options.passio and options.color_stderr is not None:
+            sys.stderr.write("\033[0m")
         if options.tray:
             tray_notify_error(str(err))
             notify_error_qubes_manager(vm.name, str(err))
@@ -73,6 +77,8 @@ def vm_run_cmd(vm, cmd, options):
     finally:
         if options.passio and options.color_output is not None:
             sys.stdout.write("\033[0m")
+        if options.passio and options.color_stderr is not None:
+            sys.stderr.write("\033[0m")
 
 def main():
     usage = "usage: %prog [options] [<vm-name>] [<cmd>]"
@@ -122,11 +128,20 @@ def main():
                       dest="color_output", default=None,
                       help="Disable marking VM output with red color")
 
+    parser.add_option("--no-color-stderr", action="store_false",
+                      dest="color_stderr", default=None,
+                      help="Disable marking VM stderr with red color")
+
     parser.add_option("--color-output", action="store", type="int",
                       dest="color_output",
                       help="Force marking VM output with given ANSI style ("
                            "use 31 for red)")
 
+    parser.add_option("--color-stderr", action="store", type="int",
+                      dest="color_stderr",
+                      help="Force marking VM stderr with given ANSI style ("
+                           "use 31 for red)")
+
     (options, args) = parser.parse_args ()
 
     if (options.passio and not options.localcmd) and options.run_on_all_running:
@@ -145,6 +160,12 @@ def main():
     elif options.color_output is False:
         options.color_output = None
 
+    if options.color_stderr is None:
+        if os.isatty(sys.stderr.fileno()) and not options.localcmd:
+            options.color_stderr = 31
+    elif options.color_stderr is False:
+        options.color_stderr = None
+
     if (options.pause or options.unpause):
         takes_cmd_argument = False
     else: