improve users plugin + users demo

Signed-off-by: Mattias Andrée <maandree@operamail.com>

1 files changed, 69 insertions, 4 deletions

diff --git a/src/plugins/users.py b/src/plugins/users.py
index a73939e..85e390c 100644
--- a/src/plugins/users.py
+++ b/src/plugins/users.py
@@ -17,6 +17,8 @@ You should have received a copy of the GNU General Public License
 along with this program.  If not, see <http://www.gnu.org/licenses/>.
 '''
 
+import os
+
 from util import *
 
 
@@ -24,19 +26,82 @@ class Users:
     '''
     Gets the logged in users and what TTY:s they are using
     
-    @variable  users:dict<str, list<str>>  Map from logged in users to acquired TTY:s
+    @variable  users:dict<str, list<str>>  Map from logged in users to acquired TTY:s (all might not be listed)
     '''
     
     
-    def __init__(self):
+    def __init__(self, method = 'who'):
         '''
         Constructor
+        
+        @param  method:str|()→itr<str>  The user listing method
         '''
-        users = spawn('w').read().decode('utf-8', 'replace').split('\n')[2 : -1]
-        users = [list(filter(lambda cell : not cell == '', line.split(' ')))[:2] for line in users]
+        users = Users.who
+        if not isinstance(method, str):
+            users = method
+        elif method == 'who':    users = Users.who
+        elif method == 'w':      users = Users.w
+        elif method == 'devfs':  users = Users.devfs
+        
+        users = users()
         self.users = {}
         for (user, tty) in users:
             if user not in self.users:
                 self.users[user] = []
             self.users[user].append(tty)
+    
+    
+    @staticmethod
+    def who():
+        '''
+        Use Coreutil's `who` command to fetch users and ttys
+        
+        @return  :list<(user:str, tty:str>)  List of user–tty pairs
+        '''
+        users = spawn('who').read().decode('utf-8', 'replace').split('\n')[:-1]
+        users = [list(filter(lambda cell : not cell == '', line.split(' ')))[:2] for line in users]
+        return users
+    
+    
+    @staticmethod
+    def w():
+        '''
+        Use procps-ng's `w` command to fetch users and ttys
+        
+        @return  :list<(user:str, tty:str>)  List of user–tty pairs
+        '''
+        users = spawn('w').read().decode('utf-8', 'replace').split('\n')[2 : -1]
+        users = [list(filter(lambda cell : not cell == '', line.split(' ')))[:2] for line in users]
+        return users
+    
+    
+    @staticmethod
+    def devfs(try_to_find_root = False):
+        '''
+        Walk /dev to find TTY acquisitions
+        
+        @param   try_to_find_root:bool       Whether to list root for /dev/tty[[:numeric:]]+ if the permission bits looks good
+        @return  :list<(user:str, tty:str>)  List of user–tty pairs
+        '''
+        import pwd
+        
+        def num(s):
+            return (not s == '') and (0 == len(list(filter(lambda c : not ('0' <= c <= '9'), s))))
+        
+        ttys = [f for f in os.listdir('/dev') if f.startswith('tty') and num(f[3:])]
+        ttys += ['pts/' + f for f in os.listdir('/dev/pts') if num(f)]
+        
+        users, rc = {}, []
+        for tty in ttys:
+            attr = os.stat('/dev/' + tty)
+            if ('/' not in tty) and (attr.st_uid == 0):
+                if not (try_to_find_root and (attr.st_mode == 0o20600)):
+                    continue
+            user = attr.st_uid
+            if user in users:
+                user = users[user]
+            else:
+                users[user] = user = pwd.getpwuid(user).pw_name
+            rc.append((user, tty))
+        return rc