StackScript Python Library

by linode
0 deployments · 0 still active · last rev. 1 year ago

Does nothing on its own. Do not deploy directly.

A collection of useful Python functions to be included in other StackScripts.

Compatible with: CentOS 7, Debian 7, Debian 8, Ubuntu 14.04 LTS
						#!/usr/bin/env python

"""
Python Library StackScript

	Author: Ricardo N Feliciano <rfeliciano@linode.com>
	Version: 1.0.0.0
	Requirements:
		- n/a

This StackScript is not meant to be directly deployed. Includes a host of
functions to do very common task on a distro, as implemented in Python. The
functions in this library are designed to be run accross the Linode Core
Distributions:
	- Ubuntu
	- CentOS
	- Debian
	- Fedora
"""


import crypt
import fcntl
import logging
import os
import platform
import pwd
import socket
import subprocess
import sys
import time
try:
	import apt
except:
	try:
		import yum
	except:
		try:
			import dnf
		except ImportError:
			print("Package manager support was not found.")
			

distro = None
"""String list: Contains details of the distribution."""


def end():
	"""End the StackScript cleanly."""
	
	# Should the StackScripts themselves be removed here at some point?
	logging.info("The StackScript has been completed.")
	subprocess.check_output('echo "The StackScript has completed. Press enter to continue." | wall -n', shell=True)


def init():
	"""Start features we consider StackScript standard."""
	
	# Sanity check for CentOS 7 & Fedora 22
	if os.path.exists("/var/log/stackscript.log"):
		sys.exit(1) # StackScript already started once, bail
	
	with open("/etc/profile.d/stackscript.sh", "w") as f:
		f.write("""#!/bin/bash
if pgrep -f "python /root/StackScript" &>/dev/null
then
	echo "####################################################################"
	echo "#####"
	echo "#####  Warning: Your StackScript is still running"
	echo "#####"
	echo "#####"
	echo "#####  Please do not make any system changes until it "
	echo "#####  completes. Log file is located at: "
	echo "#####    /var/log/stackscript.log"
	echo "####################################################################"
	echo " "
else
	echo "####################################################################"
	echo "#####"
	echo "#####  The StackScript has completed. Enjoy your system."
	echo "#####"
	echo "#####"
	echo "#####  For reference, the logfile is located at: "
	echo "#####    /var/log/stackscript.log"
	echo "####################################################################"
	echo " "
	rm /etc/profile.d/stackscript.sh
fi""")
	
	logging_start()


def logging_start(the_file="/var/log/stackscript.log", the_level=logging.INFO):
	"""Turn on logging."""
	
	logging.basicConfig(filename=the_file, level=the_level)
	logging.info("Logging has started. " + str(time.time()))


def system_detect_distro():
	"""Prepares distro information.
	
	This is critical to support the Linode Core Distributions with a single
	script.
	"""
	global distro

	# add support for logging
	
	distro = dict(zip(('distname', 'version', 'codename'),
	platform.linux_distribution(full_distribution_name=0)))
	
	distro['distname'] = distro['distname'].lower()

	if distro['distname'] in ('debian', 'ubuntu'):
		distro['family'] = "debian"
	elif distro['distname'] in ('fedora', 'centos'):
		distro['family'] = "redhat"
	else:
		raise NotImplementedError("This distribution is not supported.")


def system_IP_get():
	"""Return IPv4 address configured on eth0.
	
	This basically is a disgusting hack. Please let me know if you find a
	cleaner way to do this."""
	# add support for logging

	s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
	return socket.inet_ntoa(fcntl.ioctl(
		s.fileno(),
		0x8915, # SIOCGIFADDR
		struct.pack('256s', "eth0")
	)[20:24])


def system_IP_rdns(IP):
	"""Get PTR for given IP address."""
	# add support for logging

	return socket.gethostbyaddr(IP)[0]


def system_package_install(package, update_first=True):
	"""Install a package with the appropriate package manager."""
	# add support for logging

	if distro is None:
		system_detect_distro()

	system_update() if update_first else None
	
	if distro['family'] == "debian":
		os.environ['DEBIAN_FRONTEND'] = "noninteractive"
		cache = apt.Cache()
		pkg = cache[package]
		pkg.mark_install()
		cache.commit()
	elif distro['distname'] == "centos":
		yb = yum.YumBase()
		yb.conf.assumeyes = True
		yb.install(name=package)
		yb.resolveDeps()
		yb.processTransaction()
		yb.close()
	elif distro['distname'] == "fedora":
		dnfb = dnf.Base()
		dnfb.conf.assumeyes = True
		dnfb.read_all_repos()
		dnfb.fill_sack()
		dnfb.install(package)
		dnfb.resolve()
		dnfb.download_packages(dnfb.transaction.install_set)
		dnfb.do_transaction()
		dnfb.close()


def system_update():
	"""Uses the distro's package manager to update packages."""
	#add support for logging
	
	if distro is None:
		system_detect_distro()
	
	if distro['family'] == "debian":
		cache = apt.Cache()
		cache.update()
		cache.open(None)
		cache.upgrade()
		cache.commit()
	elif distro['distname'] == "centos":
		yb = yum.YumBase()
		yb.conf.assumeyes = True
		yb.update()
		yb.resolveDeps()
		yb.processTransaction()
		yb.close()
	elif distro['distname'] == "fedora":
		dnfb = dnf.Base()
		dnfb.conf.assumeyes = True
		#dnfb.read_all_repos() #updates were failing with this line
		dnfb.fill_sack()
		dnfb.upgrade_all()
		dnfb.resolve()
		dnfb.do_transaction()
		dnfb.close()


def user_add(username, password, groups):
	"""Creates a Linux user account.
	
	Args:
		username (String): A Linux username.
		password (String): Password for the user.
		groups (tuple): Groups that the user should be added to.
	
	Returns:
		bool: True if successful, False otherwise.
	"""

	# need to implement logging
	# need to implement group functionality

	return subprocess.call(['useradd', '-m', '-p', crypt.crypt(password, "22"), '-s', '/bin/bash', username])


def user_add_pubkey(username, key):
	"""Adds the public SSH key to the specified user."""
	# need to implement logging
	
	if username != "root":
		os.seteuid(pwd.getpwnam(username).pw_uid)
		os.setegid(pwd.getpwnam(username).pw_gid)
	
	pubkey_dir = os.path.join(os.getenv("HOME"), ".ssh")
	
	if not os.path.isdir(pubkey_dir):
		os.makedirs(pubkey_dir)
	
	with open(os.path.join(pubkey_dir, "authorized_keys")) as f:
		f.write(key)
	
	if username != "root":
		os.seteuid(0)
		os.setegid(0)