A smarter CLI – Example Shell

Here is the code of a simple shell (exampleshell) using python-configshell, libuser and hdparm.


#!/usr/bin/python

'''
Initial code by Carsten Sommer (Gonicus).
Some additions by Tom Schwaller.

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, version 3 (AGPLv3).
 
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General 
Public License along with this program.  If not, see 
lt;http://www.gnu.org/licenses/gt;.
'''

import os, sys
import configshell
import pwd, grp, libuser

#
# root config node, aka root directory of the shell, 
# there is not much in it, it only contains the 
# node "system"
#
class SystemRoot(configshell.node.ConfigNode):
    def __init__(self, shell):
        configshell.node.ConfigNode.__init__(self, '/', shell=shell)
        System(self)

#
# node "system", contains the child node "users" 
# and provides two simple commands, "uname" and "lspci"
# "lspci" can be given one optional parameter
#
class System(configshell.node.ConfigNode):
    def __init__(self, parent):
        configshell.node.ConfigNode.__init__(self, 'system', parent)
        Users(self)
        Hardware(self)  

    # simple command without any parameters
    def ui_command_uname(self):
        '''
        Displays the system uname information.
        '''
        os.system("uname -a")

    # Simple command with one optional parameter
    # it's an optional parameter if it get's initialized in the signature
    # it's an required parameter if it isn't initialized in the signature
    def ui_command_lspci(self, opt=None):
        '''
        lspci - list all PCI devices, I{opt} is an optional parameter for lspci, i.e. "-v"

        PARAMETERS
        ==========

        I{opt}
        -------
        Option for lspci.


        '''
        if opt == None:
            os.system("lspci")
        else:
            os.system("lspci %s" % opt)


#
# node "user", contains a child node for every user of the system
# and provides the "show" command, that executes "getent passwd"
#
class Users(configshell.node.ConfigNode):
    def __init__(self, parent):
        configshell.node.ConfigNode.__init__(self, 'users', parent)

        # get a list of all users and create a child node for each one
        for user in pwd.getpwall():
            User(user[0], self)
                        
    def ui_command_show(self):
        '''
        show  -  print "getent passwd" output of all users.
        '''
        os.system("getent passwd")


#
# node of a user
# provides the commands "id", "show", "chsh" and "change"
# the "chsh" and "change" commands have parameter completion
#
class User(configshell.node.ConfigNode):
    # username of the user of this node
    username = ""


    def __init__(self, username, parent):
        configshell.node.ConfigNode.__init__(self, username, parent)
        self.username = username


    # show the output of "getent passwd "
    def ui_command_ls(self):
        '''
        show  -  print "getent passwd" output of the current user.
        '''
        os.system("getent passwd %s" % self.username)


    # id command, show output of "id "
    def ui_command_id(self):
        '''
        id  -  print "id" output.
        '''
        os.system("id %s" % self.username)


    # chsh command, shell is a required parameter
    def ui_command_chsh(self, shell):
        '''
        chsh - change the shell of this user to I{shell}
            
        PARAMETERS
        ==========

        I{shell}
        -------
        Shell to set as a login shell for the current user. 


        '''
        # only execute chsh if shell is a valid login shell 
        if shell in libuser.get_user_shells():
            os.system("chsh -s %s %s" % (shell, self.username))
        else:
            self.shell.log.error("%s is not a valid login shell!" % shell)


    # parameter completion for chsh command
    def ui_complete_chsh(self, parameters, text, current_param):
        completions = []

        self.shell.log.debug("Called with params=%s, text='%s', current='%s'"
                             % (str(parameters), text, current_param))

        # if current parameter is "shell" we return completions
        # that are in /etc/shells and start with "text"
        if current_param == 'shell':
            completions = [shell for shell in libuser.get_user_shells()
                           if shell.startswith(text)]

        self.shell.log.debug("Returning completions %s." % str(completions))
        return completions


    # change command, changes various attributes of a user
    def ui_command_change(self, shell=None, group=None, home=None):
        '''
        change - changes various attributes of a user
            
        PARAMETERS
        ==========

        I{shell}
        --------
        Shell to set as a login shell for the current user. 

        I{group}
        --------
        Group that should be set as the primary group of 
        the current user and changes the group of all files of
        the users home directory.

        I{home}
        -------
        Directory that should be set as the home directory.

        SEE ALSO
        ========
        chsh

        '''
        # only execute chsh if shell is a valid login shell 
        if shell != None:
            if shell in libuser.get_user_shells():
                os.system("chsh -s %s %s" % (shell, self.username))
            else:
                self.shell.log.error("%s is not a valid login shell!" % shell)

        if home != None:
            os.system("usermod -d %s %s" % (home, self.username))

        if group != None:
            if group in [gr[0] for gr in grp.getgrall()]:
                os.system("usermod -g %s %s" % (group, self.username))
            else:
                self.shell.log.error("%s is not a valid group!" % group)


    # parameter completion for change command
    def ui_complete_change(self, parameters, text, current_param):
        completions = []

        self.shell.log.debug("Called with params=%s, text='%s', current='%s'"
                             % (str(parameters), text, current_param))

        # if current parameter is "shell" we return completions
        # that are in /etc/shells and start with "text"
        if current_param == 'shell':
            completions = [shell for shell in libuser.get_user_shells() 
                           if shell.startswith(text)]

        # if current parameter is "group" we get a list of groups
        # and return completions that start with "text"
        if current_param == 'group':
            completions = [gr[0] for gr in grp.getgrall()
                           if gr[0].startswith(text)]

        if len(completions) == 1 and not completions[0].endswith('='):
            completions = [completions[0] + ' ']


        self.shell.log.debug("Returning completions %s." % str(completions))
        return completions


class Hardware(configshell.node.ConfigNode):
    def __init__(self, parent):
        configshell.node.ConfigNode.__init__(self, 'hardware', parent)
        Partitions(self)

                        
class Partitions(configshell.node.ConfigNode):
    def __init__(self, parent):
        configshell.node.ConfigNode.__init__(self, 'partitions', parent)

        # get a list of all partitions and create a child node for each one
        for partition in self.get_partitions():
            Partition(partition, self)
                        
    def get_partitions(self):
        partitions = []
        with open('/proc/partitions', 'r') as f:
            for line in f.readlines()[2:]:
                partitions.append(line.split()[3])
        return partitions


class Partition(configshell.node.ConfigNode):
    # partition name of this node
    partition = ""

    def __init__(self, partition, parent):
        configshell.node.ConfigNode.__init__(self, partition, parent)
        self.partition = partition

    def ui_command_ls(self):
        '''
        ls  -  print "hdparm -I" output of partition.
        '''
        os.system("hdparm -I /dev/%s" % self.partition)

    def ui_command_benchmark(self):
        '''
        benchmark  -  print "hdparm -t" output of partition.
        '''
        os.system("hdparm -t /dev/%s" % self.partition)

#
# main
#
def main():
    shell = configshell.shell.ConfigShell('~/.exampleshell')
    root_node = SystemRoot(shell)
    if len(sys.argv) > 1:
        shell.run_script(sys.argv[1])
        sys.exit(0)
    else:
        shell.run_interactive()

if __name__ == "__main__":
    main()