Shell

Shell.py

# Provides a simple shell with basic
# commands for dealing with files and 
# directories.
#
# This script will try and prevent 
# unsafe file operations outside of
# Documents directory.
# 
# Add a setting named allow_unsafe and
# in the config dict and set to True
# to disable these preventative measures.

import os
import cmd
import console
import shlex

config = {}

def main():
	shell = Shell()
	shell.prompt = '> '
	shell.cmdloop()
	
colors = {'default':(0,0,0),
          'blue':(0,0,1)}

def get_color(color):
	return colors.get(color, colors['default'])

def print_in_color(color, msg):
	console.set_color(*get_color(color))
	print msg
	console.set_color(*get_color('default'))

class Shell(cmd.Cmd):

	def __init__(self):
		cmd.Cmd.__init__(self)
		self.startup_dir = os.getcwd()
				
	def safe_path(self, path):
		if config.get('allow_unsafe'):
			return True 
		paths = [self.startup_dir, os.path.normpath(path)]
		return self.startup_dir == os.path.commonprefix(paths)

	def do_ls(self, line):
		for f in sorted(os.listdir(os.getcwd())):
			if os.path.isdir(f):
				print_in_color('blue', f)
			else:
				print_in_color('default', f)

	def do_pwd(self, line):
		print os.getcwd()

	def do_cd(self, line):
		args = shlex.split(line)
		if not args:
			args = ['~/Documents']
		dir = os.path.expanduser(args[0])
		if os.path.exists(dir):
			os.chdir(dir)
			self.do_pwd(line)
		else:
			print 'Directory does not exist: ' + dir
	
	def do_mkdir(self, line):
		args = shlex.split(line)
		if not self.safe_path(os.getcwd()):
			print 'Changes to unsafe directories are disabled.'
			return 
		if os.path.exists(args[0]):
			print 'Directory already exists: ' + args[0]
			return
		os.mkdir(args[0])
		
	def do_rm(self, line):
		args = shlex.split(line)
		if not os.path.exists(args[0]):
			print 'Invalid path: ' + args[0]
			return 
		elif not self.safe_path(os.path.abspath(args[0])):
			print 'Changes to unsafe directories are disabled.'
			return 
		elif os.path.isdir(args[0]):
			os.rmdir(args[0])
		else:
			os.remove(args[0])
	
if __name__ == '__main__':
	main()