improve users plugin + users demo

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

2 files changed, 105 insertions, 4 deletions

diff --git a/examples/plugins/users b/examples/plugins/users
new file mode 100644
index 0000000..5458b37
--- /dev/null
+++ b/examples/plugins/users
@@ -0,0 +1,36 @@
+# -*- python -*-
+
+# A xpybar configuration example testing the features of plugins.xdisplay
+
+import time
+import threading
+
+from plugins.users import Users
+from plugins.clock import Clock
+
+
+OUTPUT, HEIGHT, YPOS, TOP = 0, 12, 24, True
+
+
+clock = Clock(sync_to = Clock.SECONDS)
+
+start_ = start
+def start():
+    start_()
+    async(lambda : clock.continuous_sync(lambda : bar.invalidate()))
+
+
+def redraw():
+    you = os.environ['USER']
+    def colour(user):
+        if user == 'root':     return '\033[31m%s\033[39m'
+        elif not user == you:  return '\033[33m%s\033[39m'
+        else:                  return '%s'
+    
+    users = Users(lambda : Users.devfs(True)).users
+    text = ['%s: %s' % (colour(u) % u, ' '.join(sorted(users[u]))) for u in users]
+    text = ' │ '.join(text)
+    
+    bar.clear()
+    bar.draw_coloured_text(0, 10, 0, 2, text)
+
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