Adding functions to our project Python Script

Having looked at creating simple functions in the last tutorial we are now ready to add them to our project scripts. We wish to create and delete users both from an input file and specified from the command line. To eliminate duplication of code we can call the required function to either create or delete the user account using the corresponding option. We will have three functions, createuser, deluser and openfile. the openfile function can serve across both the addfromfile and deletefromfile options.

Additionally, we can add an extra option so we can specify the expiry days of our account. We set a default to 5 days if it is not specified, but including the -e option will allow setting of the days until the account expires. We can introduce a new option to add_argument, type=int. Normally we only have strings from the command line input but this allows for the type conversion at the point of input.

#!/usr/bin/python3
import argparse, sys, os, crypt
from datetime import date, timedelta

##Create User function
def createuser(user,edays):
    user = user.lower()
    now = date.today()
    end = now + timedelta(days=edays)
    expire = end.isoformat()
    password = "Password1"
    encPassword = crypt.crypt(password,"a1")
    print("Creating user: "+ user)
    os.system("useradd -m -p " +encPassword+ " -e " +expire+" "+ user)
    for line in open("/etc/passwd"):
        if line.startswith(user + ":"):
            print(line)

##Remove users
def deluser(user):
    print ("Deleting user: " + user)
    os.system("userdel -r " + user)

##Open File function used both the add and remove users from file
def openfile(file):
    f = open(file)
    for userline in f.readlines():
        if userline == "n": #Ignore empty lines
            continue
        userline = userline.rstrip() #remove trailing newline from readlines
        if args.addfromfile:
            createuser(userline,expiredays)
        elif args.deletefromfile:
            deluser(userline)

if not os.geteuid()==0:
    sys.exit("nOnly root can create users, try sudo " +sys.argv[0]+ " n")

parser = argparse.ArgumentParser()
parser.add_argument("-e","--expire",type=int,help="Days to expire account, default 5 days")
parser.add_argument("-a","--add", nargs="+", help="Creates local Linux Accounts")
parser.add_argument("-d","--delete",nargs="+", help="Removes local Linux Accounts")
parser.add_argument("-f","--addfromfile",help="Create Local Linux Accounts from file")
parser.add_argument("-r","--deletefromfile",help="Delete Local Linux Accounts from file")
args = parser.parse_args()

if args.expire:
    expiredays = args.expire
else:
    expiredays = 5

if args.add:
    for u in args.add:
        createuser(u, expiredays)
elif args.delete:
    for u in args.delete:
        deluser(u)

elif args.addfromfile:
    openfile(args.addfromfile)
elif args.deletefromfile:
    openfile(args.deletefromfile)