Added hooks support

4 files changed, 99 insertions, 5 deletions

diff --git a/machammer/decorators.py b/machammer/decorators.py
new file mode 100644
index 0000000..32ec511
--- /dev/null
+++ b/machammer/decorators.py
@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+
+import re
+import os
+import stat
+import inspect
+
+import hooks
+
+
+def login(func):
+    def func_wrapper(hook='login'):
+        path = '/var/root/Library/mh_%shook.py' % hook
+
+        # skip the decorator and function def
+        s = inspect.getsource(func).split('\n')[2:]
+
+        # determine indent level for re-indentation
+        indent = re.match(r'^(\s+)', s[0]).group(0)
+        f = open(path, 'w')
+        f.write('#!/usr/bin/env python\n')
+
+        for l in s:
+            f.write(l.replace(indent, '') + '\n')
+
+        f.close()
+        # only root should read and execute
+        os.chown(path, 0, 0)
+        os.chmod(path, stat.S_IXUSR | stat.S_IRUSR)
+
+        if hook == 'login':
+            hooks.login(path)
+        else:
+            hooks.logout(path)
+
+        return path
+
+    return func_wrapper
diff --git a/machammer/defaults.py b/machammer/defaults.py
new file mode 100644
index 0000000..9674a8e
--- /dev/null
+++ b/machammer/defaults.py
@@ -0,0 +1,24 @@
+# -*- coding: utf-8 -*-
+
+from .functions import call, check_output
+
+DEFAULTS_PATH = '/usr/bin/defaults'
+
+
+def defaults(*args):
+    if any(i == 'read' for i in args):
+        return check_output(DEFAULTS_PATH, *args)
+
+    return call(DEFAULTS_PATH, *args)
+
+
+def get(*args):
+    return defaults('read', *args)
+
+
+def set(*args):
+    return defaults('write', *args)
+
+
+def delete(*args):
+    return defaults('delete', *args)
diff --git a/machammer/functions.py b/machammer/functions.py
index 900569d..1a7b1d1 100644
--- a/machammer/functions.py
+++ b/machammer/functions.py
@@ -1,9 +1,9 @@
 # -*- coding: utf-8 -*-
 
 import os
+import sys
 import plistlib
 import subprocess
-import sys
 import tempfile
 
 from xml.parsers.expat import ExpatError
@@ -78,7 +78,7 @@ def exec_jar(path, user):
     if not os.path.exists(javapath):
         raise ValueError('Looks like your machine does not have Java installed')
 
-    subprocess.call(['/bin/launchctl', 'asuser', user, javapath, '-jar', path, '-silent'])
+    call('/bin/launchctl', 'asuser', user, javapath, '-jar', path, '-silent')
 
 
 def osascript(s):
@@ -146,7 +146,7 @@ def mount_image(dmg):
     except ExpatError:  # probably a EULA-image, return None instead of breaking
         return None
 
-    for p in [p.get('mount-point') for p in plist.get('system-entities')]:
+    for p in [x.get('mount-point') for x in plist.get('system-entities')]:
         if p and os.path.exists(p):
             return p
 
@@ -173,7 +173,8 @@ def mount_afp(url, username, password, mountpoint=None):
     """Mount AFP share."""
     if mountpoint is None:
         mountpoint = tempfile.mkdtemp()
-    subprocess.call(['/sbin/mount_afp', 'afp://%s:%s@%s' % (username, password, url), mountpoint])
+    url = 'afp://%s:%s@%s' % (username, password, url)
+    call('/sbin/mount_afp', url, mountpoint)
     return mountpoint
 
 
@@ -183,7 +184,7 @@ def umount(path):
 
 
 def install_su(restart=True):
-    """Install all available Apple software Updates, restart if update requires it."""
+    """Install all available Apple software Updates, restart if any update requires it."""
     su_results = subprocess.check_output(['/usr/sbin/softwareupdate', '-ia'])
     if restart and ('restart' in su_results):
         tell_app('Finder', 'restart')
diff --git a/machammer/hooks.py b/machammer/hooks.py
new file mode 100644
index 0000000..8826ad9
--- /dev/null
+++ b/machammer/hooks.py
@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+
+import defaults
+
+PREF_DOMAIN = 'com.apple.loginwindow'
+
+
+def login(path=None):
+    """Set login hook to path, or disable login hook."""
+    if path is None:
+        return defaults.delete(PREF_DOMAIN, 'LoginHook')
+
+    return defaults.set(PREF_DOMAIN, 'LoginHook', path)
+
+
+def logout(path=None):
+    """Set logout hook to path, or disable logout hook."""
+    if path is None:
+        return defaults.delete(PREF_DOMAIN, 'LogoutHook')
+
+    return defaults.set(PREF_DOMAIN, 'LogoutHook', path)
+
+
+def reboot(path=None):
+    """Set reboot hook to path, or disable reboot hook."""
+    pass
+
+
+def shutdown(path=None):
+    """Set shutdown hook to path, or disable shutdown hook."""
+    pass