#!/usr/bin/env python
# by Hoong Ern Ng
# hoongern@gmail.com
# 28 August 2011

from subprocess import Popen, PIPE
from Tkinter import * 
import tkMessageBox, os, sys

# Path to various binaries
pmset = '/usr/bin/pmset'
ps = '/bin/ps'

# Use ps to get a process name, given a pid
def getpname(pid):
	try:
		pname = Popen([ps, '-p', pid, '-o', 'comm='], stdout=PIPE, stderr=PIPE).communicate()[0].strip()
	except OSError as (errno, strerror):
		Msg = "Unable to run '" + ps + "'! (Error " + str(errno) + ": " + strerror + ")"
		PrintMsg(Msg, DoExit = True, Err = True)

	return pname

# Use applescript to sleep the system
def DoSleep():
	pid = Popen([pmset, 'sleepnow']).pid
	sys.exit()

# Print a message using tkMessageBox
def PrintMsg(Msg, DoExit = False, Err = False):
	if Err:
		print "Error: ",
	tkMessageBox.showinfo("Sleep Check", Msg, icon='info')
	if DoExit:
		sys.exit()

# Returns the description of a numeric hibernate mode
def HibernateStr(n):
	if n == -1:
		return "unknown sleep/hibernation mode"
	elif n == 0:
		return "sleep"
	elif n == 1:
		return "hibernation"
	elif n == 3:
		return "smart sleep"
	elif n == 5:
		return "secure hibernation"
	elif n == 7:
		return "secure smart sleep"
	else:
		return "hibernation mode '" + n + "'"
		
root = Tk()
root.withdraw()

# Run pmset -g and get output
try:
	pmset_g_output = Popen([pmset, '-g'], stdout=PIPE, stderr=PIPE).communicate()[0].splitlines()
except OSError as (errno, strerror):
	Msg = "Unable to run '" + pmset + "'! (Error " + str(errno) + ": " + strerror + ")"
	PrintMsg(Msg, DoExit = True, Err = True)
	
# Parse pmset output for hibernate mode
hibernate_mode = -1
for line in pmset_g_output:
	if line.strip().startswith('hibernatemode'):
		hibernate_mode = int(line.strip()[-1:])
		
# Get pmset -g assertions
try:
	pmset_ga_output = Popen([pmset, '-g', 'assertions'], stdout=PIPE, stderr=PIPE).communicate()[0].splitlines()
except OSError as (errno, strerror):
	Msg = "Unable to run '" + pmset + "'! (Error " + str(errno) + ": " + strerror + ")"
	PrintMsg(Msg, DoExit = True, Err = True)

# Parse assertion output for sleep status
num_assertions = 0
for line in pmset_ga_output:
    if line.strip().startswith('PreventUserIdleSystemSleep'):
    	n = line[line.rfind(" ")+1:]
    	num_assertions = num_assertions + int(n)
    if line.strip().startswith('PreventSystemSleep'):
    	n = line[line.rfind(" ")+1:]
    	num_assertions = num_assertions + int(n)

# Sleep if no assertions
if num_assertions == 0:
   	Msg = "Enter " + HibernateStr(hibernate_mode) + " mode now?"
   	if tkMessageBox.askquestion("Sleep Check", Msg, default='yes') == "yes": 
   		DoSleep()
   	else:
   		sys.exit()

# We have conditions preventing sleep
assertions = []
append = False
for line in pmset_ga_output:
	if line.strip().startswith('Listed by owning process:'):
		append = True
	if append and line == '':
		break
	if line.strip().startswith('pid'):
		if line.find('PreventSystemSleep') != -1 or line.find('PreventUserIdleSystemSleep') != -1:
			assertions.append(line.strip())

# Parse conditions into a nice "user friendly" list
conditions = []
for entry in assertions:
	if entry.find('com.apple.AppleFileServer') != -1:
		conditions.append('A remote user is connected via File Sharing. Disconnect the user or disable file sharing.')
	elif entry.find('com.apple.audio') != -1:
		conditions.append('An application is playing audio.')
	elif entry.find('com.apple.InternetSharing') != -1:
		conditions.append('Internet Sharing is enabled.')
	elif entry.find('com.apple.powermanagement.ttyassertion') != -1:
		conditions.append('A remote user is connected to a tty (i.e. via telnet/SSH). Disconnect the user or run "sudo pmset -a ttyskeepawake 0" to allow system sleep even when remote users are connected.')
	else:
		# Get the process name
		pid = entry[4:entry.find(':')]
		pname = getpname(pid)
		condition_name = entry.split('named: ')[1]
		conditions.append('"' + pname + '" is preventing system sleep. (' + condition_name + ')')

# Print out message of conditions
if len(conditions) == 1:
	msgpart = " condition is "
else:
	msgpart = " conditions are "
Msg = str(len(conditions)) + msgpart + 'preventing system sleep/hibernation:\n\n - ' + '\n\n - '.join(conditions)

PrintMsg(Msg, DoExit=True)