Script to setup an Ubuntu machine to my liking

ubuntuSetup.py

DESCRIPTION=\
'''
Brief:
    ubuntuSetup.py - A script to setup an Ubuntu-family system to the way I like it.
        Tested on Ubuntu 16.04 LTS. Should work on Ubuntu 14.04 LTS.
            Enabled on Debian but not tested at all... may work, may not.

Description:
    Some of the things this installs and configures:
        - VS Code
        - xrdp
        - Virtualbox
        - Vino (no encryption)
        - Atom
        - Chromium
        - Mercurial with Tortoise Hg
        - Python pip
        - Python tab-completion
        - Openssh Server (with x11 forwarding)
        - No password on sudo
        - Samba
        - Screen
        - Git with Git Extensions

        - Proxy (Optional via cmd line switches)

Author(s):
    Charles Machalow
'''
import argparse
import platform
import os
import subprocess, sys

global DEBUG
DEBUG = 2

EXIT_FAILURE = -1
EXIT_SUCCESS = 0
DESKTOP_ICON = \
'''#!/usr/bin/env xdg-open
[Desktop Entry]
Version=1.0
Type=Application
Terminal=%s
Exec=%s
Name=%s
'''
DISABLE_AUTO_SUSPEND = \
'''
[Disable suspend (upower)]
Identity=unix-user:*
Action=org.freedesktop.upower.suspend
ResultActive=no
ResultInactive=no
ResultAny=no

[Disable suspend (logind)]
Identity=unix-user:*
Action=org.freedesktop.login1.suspend
ResultActive=yes
ResultInactive=no
ResultAny=yes

[Disable suspend when others are logged in (logind)]
Identity=unix-user:*
Action=org.freedesktop.login1.suspend-multiple-sessions
ResultActive=yes
ResultInactive=no
ResultAny=yes
'''
GIT_EXT_SH = \
'''
#!/bin/bash
mono /usr/local/bin/GitExtensions/GitExtensions.exe "$@" &
'''
NO_PASSWORD_SUDO = "%s ALL=(ALL) NOPASSWD:ALL"
PYTHON_TAB_COMPLETE = \
"""
try:
    import readline
    import rlcompleter
    readline.parse_and_bind("tab: complete")
except:
    pass
"""
SUPPORTED_DISTROS = ['UBUNTU', 'DEBIAN']
VINO_ENABLE = 'gsettings set org.gnome.Vino enabled true'
VINO_ENCRYPTION_FIX = 'gsettings set org.gnome.Vino require-encryption false'
VINO_PROMPT = 'gsettings set org.gnome.Vino prompt-enabled false'
VINO_START = '/usr/lib/vino/vino-server > /dev/null 2>&1 &'
VINO_VIEW = 'gconftool-2 -s -t bool /desktop/gnome/remote_access/view_only false'
VINO_LINE = VINO_VIEW + "\n" + VINO_ENABLE + "\n" + VINO_START + "\n" + VINO_ENCRYPTION_FIX + "\n" + VINO_PROMPT +"\n"

def LOG_DEBUG(s):
    '''
    Debug printing
    '''
    global DEBUG
    if DEBUG >= 3:
        sys.stderr.write(s + "\n")

def LOG_INFO(s):
    '''
    Info printing
    '''
    global DEBUG
    if DEBUG >= 2:
        sys.stderr.write(s + "\n")

def LOG_ERROR(s):
    '''
    Error printing
    '''
    global DEBUG
    if DEBUG >= 1:
        sys.stderr.write(s + "\n")

def isDebug():
    '''
    Returns True if a higher level of debug besides silent/normal is set
    '''
    global DEBUG
    return DEBUG > 2

def failAndExit(s):
    '''
    Causes the entire script to fail and exit
    '''
    sys.stderr.write('%s\n' % s)
    os._exit(EXIT_FAILURE)

def isSupportedDistro():
    '''
    Returns True if this is a supported distro
    '''
    dist, ver, codeName = platform.linux_distribution()
    LOG_DEBUG("%s %s (%s)" % (dist, ver, codeName))
    if dist.upper() in SUPPORTED_DISTROS:
        return True

    LOG_ERROR("Unsupported Distro: %s" % dist)
    return False

def addIfNotInFile(fileName, txtToAdd):
    '''
    Adds the text to the file if it isn't already there
    '''
    try:
        with open(fileName, 'a+') as f:
            text = f.read()
            if txtToAdd in text:
                LOG_DEBUG("Not adding %s to %s as it already has that text" % (txtToAdd, fileName))
            else:
                f.write(txtToAdd)
                LOG_DEBUG("Added %s to %s" % (txtToAdd, fileName))
        return True
    except:
        LOG_ERROR("Unable to create: %s" % fileName)
        return False

def getUserName():
    '''
    Returns the username
    '''
    return os.environ['SUDO_USER']

def getArgs():
    '''
    Returns the result of parsing args
    '''
    parser = argparse.ArgumentParser(description=DESCRIPTION, formatter_class=argparse.RawDescriptionHelpFormatter)
    verbosity = parser.add_mutually_exclusive_group()
    verbosity.add_argument("-d", "--debug", help="Print some debug info during the process",
        action="store_true")
    verbosity.add_argument("-s", "--silent", help="Don't even print most errors'",
        action="store_true")
    parser.add_argument("--http_proxy", help="If given, an http proxy to set for environment variables, wget, hg, git",
        type=str)
    parser.add_argument("--https_proxy", help="If given, an https proxy to set for environment variables, wget, hg, git",
        type=str)
    parser.add_argument("--hostname", help="If given, will set this system's hostname to the given string",
        type=str)
    args = parser.parse_args()
    return args

def callSystemCommand(command):
    '''
    Calls a system command with respect to DEBUG and returns True on success
    '''
    LOG_INFO("Calling bash command: " + command)

    if isDebug():
        retCode = os.system(command) # Do it like this to print output in realtime
    else:
        try:
            output = subprocess.check_output(command, stderr=subprocess.STDOUT, shell=True)
            LOG_DEBUG(output)
            retCode = EXIT_SUCCESS
        except subprocess.CalledProcessError as ex:
            LOG_DEBUG(ex.output)
            LOG_ERROR("Failed to call: %s" % command)
            retCode = ex.returncode

    return retCode == EXIT_SUCCESS

def callAptGet(s):
    '''
    Calls sudo apt-get %s -y

    Returns True on success, False on failure
    '''
    scriptCall = "sudo apt-get %s -y" % s
    LOG_DEBUG("About to call %s" % scriptCall)

    return callSystemCommand(scriptCall)

def aptGetUpdateUpgrade():
    '''
    Calls apt-get update, install -f, and then apt-get upgrade
    '''
    if not callAptGet("update"):
        failAndExit("Could not update apt-get")

    if not callAptGet("install -f"):
        failAndExit("Could not intall -f apt-get")

    if not callAptGet("upgrade"):
        failAndExit("Could not upgrade apt-get")

def isRoot():
    '''
    Returns True if we are root
    '''
    return 'SUDO_USER' in os.environ

def noPasswordOnSudo():
    '''
    Makes it so the current user doesn't need password on sudo'
    '''
    line = NO_PASSWORD_SUDO % getUserName()

    LOG_DEBUG("About to make it so we don't need the password on sudo")

    with open('/etc/sudoers', 'r') as f:
        file = f.read()

    if line in file:
        LOG_DEBUG("Sudoers file already has this feature. Skipping.")
    else:
        callSystemCommand('''sudo bash -c 'echo "%s" | (EDITOR="tee -a" visudo)' ''' % line)
        LOG_DEBUG("Done.")

def fixVinoEncryptionAndEnable():
    '''
    Disables vino encryption to allow most clients to connect.
    Also adds the fix to /etc/rc.local to be applied on startup
    '''
    callAptGet('install vino')

    if not os.path.isfile('/etc/rc.local'):
        callSystemCommand('sudo touch /etc/rc.local')
        callSystemCommand('sudo chmod +x /etc/rc.local')
        LOG_DEBUG("Created /etc/rc.local file")

    with open('/etc/rc.local', 'r') as f:
        file = f.read()

    if VINO_LINE in file:
        LOG_DEBUG("/etc/rc.local already has the vino fix. Skipping")
    else:
        with open('/etc/rc.local', 'w') as f:
            f.write(file + "\n" + VINO_LINE)

    callSystemCommand('sudo pkill vino || true') # kill vino if running. Don't error if not running
    callSystemCommand(VINO_LINE)
    LOG_DEBUG("Done.")

def isInstalled(name):
    '''
    Returns True if the passed in item is installed
    '''
    # Don't print failed to calls for this
    LOG_INFO("Checking if %s is installed" % name)
    global DEBUG
    debugBack = DEBUG
    DEBUG = 0
    retVal = callSystemCommand('sudo dpkg -s %s' % name)
    DEBUG = debugBack
    return retVal

def installAtom():
    '''
    Installs the atom text editor
    '''
    if not isInstalled('atom'):
        if callSystemCommand('wget https://github.com/atom/atom/releases/download/v1.11.2/atom-amd64.deb -O /tmp/atom.deb'):
            LOG_DEBUG("Successfully used wget to get the Atom.deb file")
            if callSystemCommand('sudo dpkg -i /tmp/atom.deb'):
                LOG_DEBUG("Finished installing Atom.deb")
                return True

        LOG_ERROR("Failed to wget Atom")
        return False
    else:
        LOG_DEBUG("Atom is already on this system. Skipping.")

def installVsCode():
    '''
    Installs the VS Code editor
    '''
    if not isInstalled('code'):
        if callSystemCommand('wget https://go.microsoft.com/fwlink/?LinkID=760868 -O /tmp/vscode.deb'):
            LOG_DEBUG("Successfully used wget to get the vscode.deb file")
            if callSystemCommand('sudo dpkg -i /tmp/vscode.deb'):
                LOG_DEBUG("Finished installing vscode.deb")
                return True

        LOG_ERROR("Failed to wget VSCode")
        return False
    else:
        LOG_DEBUG("VSCode is already on this system. Skipping.")

def setupProxy(http, https):
    '''
    Sets the proxy for various applications
    '''
    if http and https:
        globalProxyLine = "\nHTTP_PROXY=%s\nHTTPS_PROXY=%s\n" % (http, https)
        LOG_DEBUG("Setting up http and https proxies")
    elif http:
        globalProxyLine = "\nHTTP_PROXY=%s\n" % (http)
        LOG_DEBUG("Setting up http proxy")
    elif https:
        globalProxyLine = "\nHTTPS_PROXY=%s\n" % (https)
        LOG_DEBUG("Setting up https proxy")
    else:
        LOG_DEBUG("Skipping proxy setup")
        return True # neither

    # Global via environment variables
    with open('/etc/environment', 'r') as f:
        fileText = f.read()

    if globalProxyLine not in fileText:
        fileText += globalProxyLine
        with open('/etc/environment', 'w') as f:
            f.write(fileText)
        LOG_DEBUG("Added http/https proxy environment variables to /etc/environment")
    else:
        LOG_DEBUG("Skipping adding to /etc/environment as the variables are already there")

    # wget
    wgetrcLocation = "/home/%s/.wgetrc" % getUserName()
    if os.path.exists(wgetrcLocation):
        with open(wgetrcLocation, 'r') as f:
            wgetrc = f.read()
    else:
        wgetrc = ""

    wgetRcProxyLine = "\nuse_proxy=yes\n%s" % globalProxyLine.lower()
    if wgetRcProxyLine not in wgetrc:
        wgetrc += wgetRcProxyLine
        with open(wgetrcLocation, 'w') as f:
            f.write(wgetRcProxyLine)
        LOG_DEBUG("Added proxy info to .wgetrc file")
    else:
        LOG_DEBUG("Skipping adding proxy info to .wgetrc as it is already there")

    # git
    if http:
        callSystemCommand("git config --global http.proxy $HTTP_PROXY")

    if https:
        callSystemCommand("git config --global https.proxy $HTTPS_PROXY")

    LOG_DEBUG("Set proxy for git")
    LOG_DEBUG("Done setting proxies")

def addDesktopIcon(name, command, terminal=False):
    '''
    Adds a desktop icon with this name to do the given command
    Returns True if the icon was added
    '''

    desktopIcon = DESKTOP_ICON % (str(terminal).lower(), command, name)

    try: # make sure Desktop directory exists
        os.makedirs(r'/home/%s/Desktop/' % getUserName())
    except:
        pass

    fileLoc = r'/home/%s/Desktop/%s.desktop' % (getUserName(), name)
    if os.path.exists(fileLoc):
        LOG_DEBUG("File exists. Not overwritting (%s)" % fileLoc)
        return False

    with open(fileLoc, 'w') as f:
        f.write(desktopIcon)

    callSystemCommand('sudo chown %s %s' % (getUserName(), fileLoc))
    callSystemCommand('sudo chmod +x %s' % fileLoc)

    LOG_DEBUG("Icon created: %s" % fileLoc)
    return True

def addPyTabCompletion():
    '''
    Adds things to .pythonrc to get readline and tab complete working
    '''
    pyrc = '/home/%s/.pythonrc' % getUserName()

    with open('/etc/environment', 'a+') as f:
        text = f.read()
        if 'PYTHONSTARTUP' not in text:
            LOG_DEBUG("Added PYTHONSTARTUP environment variable")
            f.write("\nPYTHONSTARTUP=%s\n" % pyrc)
        else:
            LOG_DEBUG("Already have PYTHONSTARTUP environment variable. Skipping")

    if os.path.exists(pyrc):
        with open(pyrc, 'r') as f:
            fileText = f.read()
    else:
        fileText = ""

    if PYTHON_TAB_COMPLETE in fileText:
        LOG_DEBUG("Python tab complete code already exists. Skipping.")
        return False

    with open(pyrc, 'a') as f:
        f.write("\n%s" % PYTHON_TAB_COMPLETE)

    LOG_DEBUG("Done adding Python tab complete")
    return True

def getOsCodename():
    '''
    Returns the OS codename
    '''
    out = subprocess.check_output("lsb_release -c", shell=True).decode()
    return out.split(':')[-1].strip()

def installVirtualBox():
    '''
    Installs VirtualBox
    '''
    if (isInstalled('virtualbox') and isInstalled('virtualbox-qt')) is False:
        if callSystemCommand("""echo "deb http://download.virtualbox.org/virtualbox/debian %s contrib" >> /etc/apt/sources.list""" % getOsCodename()):
            LOG_DEBUG("Added virtualbox apt repository")

        if callSystemCommand("wget -q https://www.virtualbox.org/download/oracle_vbox_2016.asc -O- | sudo apt-key add -") and \
            callSystemCommand("wget -q https://www.virtualbox.org/download/oracle_vbox.asc -O- | sudo apt-key add -"):
            LOG_DEBUG("Successfully added keys")

        aptGetUpdateUpgrade()
        if callAptGet("install virtualbox") and callAptGet("install virtualbox-qt"):
            LOG_DEBUG("Successfully installed virtualbox")
            return True

        LOG_ERROR("Unable to install virtualbox")
        return False

    else:
        LOG_DEBUG("Not installing virtualbox as it is already installed. Still will try virtualbox-qt as sometimes it is missing.")
        callAptGet("install virtualbox-qt")
        return True

def disableScreensaverAndLock():
    '''
    Disables the screensaver and lock
    '''

    if not addIfNotInFile('/etc/polkit-1/localauthority/50-local.d/com.ubuntu.disable-suspend.pkla', DISABLE_AUTO_SUSPEND):
        LOG_ERROR("Couldn't add ubuntu disable suspend file")

    callSystemCommand('sudo gsettings set org.gnome.desktop.screensaver lock-enabled false')
    callSystemCommand('sudo gsettings set org.gnome.desktop.session idle-delay 0')
    callSystemCommand('sudo xset dpms force off')
    callSystemCommand('''sudo gsettings set org.gnome.settings-daemon.plugins.power sleep-inactive-ac-timeout '0' && sudo gsettings set org.gnome.settings-daemon.plugins.power sleep-inactive-battery-timeout '0' ''')
    callSystemCommand("sudo xset s off")

def setupXrdp():
    '''
    Installs xrdp and an alternate window manager
    '''
    if not callAptGet("install xrdp"):
        failAndExit("Couldn't install xrdp")

    if not callAptGet("install xfce4-session"):
        failAndExit("Couldn't install xfce4-session")

    if not callAptGet("install xfce4-panel"):
        failAndExit("Couldn't install xfce4-panel")

    if not callAptGet("install thunar"):
        failAndExit("Couldn't install thunar")

    callSystemCommand('mkdir /home/%s/.config/xfce4 || true' % getUserName())
    addIfNotInFile('/home/%s/.config/xfce4/helpers.rc' % getUserName(), '\nFileManager=Thunar\n') # set thunar as default to avoid nautilus taking over the desktop
    callSystemCommand('sudo chown %s /home/%s/.config -R' % (getUserName(), getUserName())  )
    callSystemCommand('''sudo sed -i 's/BIG-REQUESTS/_IG-REQUESTS/' /usr/lib/x86_64-linux-gnu/libxcb.so.1''') # otherwise vscode and atom wont open in xrdp
    addIfNotInFile('/home/%s/.xsession' % getUserName(), "\nxfce4-session\n")
    callSystemCommand('sudo service xrdp restart')

    LOG_DEBUG("xrdp is fully setup")

def installHg():
    '''
    Installs hg and tortoisehg
    '''
    if not callAptGet("install mercurial"):
        LOG_ERROR("Couldn't install mercurial")
        return False

    callSystemCommand('sudo add-apt-repository ppa:tortoisehg-ppa/releases -y')
    aptGetUpdateUpgrade()
    callAptGet('install tortoisehg')
    callAptGet('tortoisehg-nautilus') # for workbench right click in nautilus

def installGitAndGitExtensions():
    '''
    Installs git and gitextensions
    '''
    LOG_DEBUG("Starting git installation")
    if not callAptGet("install git"):
        failAndExit("Couldn't install git")

    callAptGet("install kdiff3")

    callSystemCommand("sudo apt-key adv --keyserver hkp://keyserver.ubuntu.com:80 --recv-keys 3FA7E0328081BFF6A14DA29AA6A19B38D3D831EF")
    callSystemCommand('echo "deb http://download.mono-project.com/repo/debian wheezy main" | sudo tee /etc/apt/sources.list.d/mono-xamarin.list')
    aptGetUpdateUpgrade()
    callAptGet("install mono-complete")

    callSystemCommand("wget https://github.com/gitextensions/gitextensions/releases/download/v2.51.01/GitExtensions-2.51.01-Mono.zip -O /tmp/gitext.zip")
    callSystemCommand("unzip -o /tmp/gitext.zip -d /usr/local/bin/")
    with open('/usr/local/bin/gitext', 'w') as f:
        f.write(GIT_EXT_SH)
    callSystemCommand('chmod +x /usr/local/bin/gitext')
    LOG_DEBUG("Completed git extensions install")

def updateHostname(newHost):
    '''
    In-place updates the system's hostname to the given one
    '''
    oldHost = subprocess.check_output("hostname", shell=True).strip()
    newHost = newHost.replace(" ", "_")
    callSystemCommand("sudo hostname %s" % newHost)
    callSystemCommand(r"sudo sed -i -e 's/%s/%s/g' /etc/hostname" % (oldHost, newHost))
    callSystemCommand(r"sudo sed -i -e 's/%s/%s/g' /etc/hosts" % (oldHost, newHost))
    return True

def runSetup(myArgs=None):
    '''
    The equivalent of a 'main' method for this script
    '''
    global DEBUG

    if not isSupportedDistro():
        failAndExit("Will not continue on unsupported distro")

    if myArgs is not None:
        # Set verbosity
        if myArgs.debug:
            DEBUG = 3
        if myArgs.silent:
            DEBUG = 0

        # hostname
        if myArgs.hostname:
            updateHostname(myArgs.hostname)

    if not isRoot():
        failAndExit("This must be run as root")

    # do the proxy setup twice, since we may need it to get git, but then again after git
    setupProxy(myArgs.http_proxy, myArgs.https_proxy)
    installGitAndGitExtensions()
    setupProxy(myArgs.http_proxy, myArgs.https_proxy)

    aptGetUpdateUpgrade()

    noPasswordOnSudo()

    disableScreensaverAndLock()

    fixVinoEncryptionAndEnable()

    if not callAptGet("install perl"):
        LOG_ERROR("Couldn't install perl")

    if not callAptGet("install python-pip"):
        LOG_ERROR("Couldn't install pip")

    if not callAptGet("install python"):
        LOG_ERROR("Couldn't install python")

    if not callAptGet("install python-dev"):
        LOG_ERROR("Couldn't install python-dev")

    if not callAptGet("install python3-pip"):
        LOG_ERROR("Couldn't install pip3")

    if not callAptGet("install python3"):
        LOG_ERROR("Couldn't install python3")

    if not callAptGet("install python3-dev"):
        LOG_ERROR("Couldn't install python3-dev")

    if not callAptGet("install openssh-server"):
        failAndExit("Couldn't install openssh-server")
    else:
        addIfNotInFile('/etc/ssh/sshd_config', '\nX11Forwarding yes\n')
        callSystemCommand('sudo service sshd restart')
        callAptGet("install sshpass")

    if not callAptGet("install htop"):
        LOG_ERROR("Couldn't install htop")

    if not callAptGet("install xauth"):
        LOG_ERROR("Couldn't install xauth")

    if not callAptGet("install ffmpeg"):
        LOG_ERROR("Couldn't install ffmpeg")

    if not callAptGet("install mplayer"):
        LOG_ERROR("Couldn't install mplayer")

    if not callAptGet("install gparted"):
        LOG_ERROR("Couldn't install gparted")

    if not callAptGet("install gnome-session-flashback"):
        if not callAptGet("install gnome-session-fallback"):
            LOG_ERROR("Couldn't install gnome-session-flashback/fallback")

    if not callAptGet("install sg3-utils"):
        LOG_ERROR("Could not install sg3-utils")

    if not callAptGet("install blktrace"):
        LOG_ERROR("Could not install blktrace")

    if not callAptGet("install sysstat"):
        LOG_ERROR("Could not install sysstat")

    if not callAptGet("install fio"):
        LOG_ERROR("Could not install fio")

    if not callAptGet("install linux-tools-common"):
        LOG_ERROR("Could not install linux-tools-common")

    setupXrdp()

    if not callAptGet("install chromium-browser"):
        LOG_ERROR("Couldn't install chromium-browser")
    else:
        addDesktopIcon('chromium', 'chromium-browser')

    if not callAptGet("install samba"):
        LOG_ERROR("Couldn't install samba")

    if not callAptGet("install screen"):
        LOG_ERROR("Couldn't install screen")

    if not callAptGet("install nload"):
        LOG_ERROR("Couldn't install nload")
    else:
        addDesktopIcon('nload', 'nload', True)

    if not callSystemCommand("sudo pip install pyreadline"):
        LOG_ERROR("Couldn't install pyreadline")

    addPyTabCompletion()

    installAtom()
    addDesktopIcon('Atom', 'atom')

    installVsCode()
    addDesktopIcon('VSCode', 'code')

    installVirtualBox()
    addDesktopIcon('VirtualBox', 'virtualbox')

    disableScreensaverAndLock()

    LOG_INFO("Finished!")

if __name__ == '__main__':
    runSetup(getArgs())