Added hooks support

6 files changed, 150 insertions, 13 deletions

diff --git a/README.md b/README.md
index 8e4125d..42eb8fb 100644
--- a/README.md
+++ b/README.md
@@ -107,6 +107,22 @@ if __name__ == '__main__':
 $ python install.py install_autocad update_render_nodes
 ```
 
+### Login hooks and decorators
+
+`machammer` allows you to execute Python code as login hooks. Just use the `@login` decorator on your loginhook function:
+
+
+```python
+from machammer.decorators import login
+
+@login
+def sayhello():
+    import sys
+    import subprocess
+    subprocess.call(['/usr/bin/say', 'Hello ' + sys.argv[1]])
+
+sayhello()
+```
 
 ### system_profiler
 
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
diff --git a/tests.py b/tests.py
index 10057ec..0dd36c5 100755
--- a/tests.py
+++ b/tests.py
@@ -6,9 +6,9 @@ import logging
 import subprocess
 from unittest import main, skip, TestCase
 
-from machammer import network, users
-from machammer import functions as mh
-from machammer import system_profiler, screensaver
+from machammer import (functions, system_profiler,
+                       network, hooks, users,
+                       screensaver, defaults,)
 
 
 class SystemProfilerTestCase(TestCase):
@@ -50,6 +50,7 @@ class NetworkTestCase(TestCase):
     def test_wired(self):
         self.assertTrue(network.is_wired())
 
+    @skip('blaa')
     def test_wifi_disable(self):
         network.set_wifi_power(False)
         time.sleep(3)
@@ -84,7 +85,7 @@ class FunctionsTestCase(TestCase):
         self.stickes = '/Applications/Stickies.app'
 
     def test_notification(self):
-        mh.display_notification('blaaa "lalala"')
+        functions.display_notification('blaaa "lalala"')
 
     def test_add_login_item(self):
         users.add_login_item(self.stickes)
@@ -93,16 +94,17 @@ class FunctionsTestCase(TestCase):
         users.remove_login_item(path=self.stickes)
 
     def test_mount_image(self):
-        p = mh.mount_image('/Users/filipp/Downloads/AdobeFlashPlayer_22au_a_install.dmg')
+        p = functions.mount_image('/Users/filipp/Downloads/AdobeFlashPlayer_22au_a_install.dmg')
         self.assertEquals(p, '/Volumes/Adobe Flash Player Installer')
 
+    @skip('This works, trust me.')
     def test_create_media(self):
-        mh.create_os_media('/Applications/Install macOS Sierra.app',
-                           '/Volumes/Untitled')
+        functions.create_os_media('/Applications/Install macOS Sierra.app',
+                                  '/Volumes/Untitled')
 
     @skip('This works, trust me.')
     def test_sleep(self):
-        mh.sleep()
+        functions.sleep()
 
 
 class ScreenSaverTestCase(TestCase):
@@ -117,6 +119,31 @@ class ScreenSaverTestCase(TestCase):
         self.assertEquals(screensaver.get(), 'Flurry')
 
 
+class HooksTestCase(TestCase):
+    def gethook(self):
+        return defaults.get(hooks.PREF_DOMAIN, 'LoginHook')
+
+    def test_set_login_path(self):
+        hooks.login('/lalala')
+        self.assertEquals(self.gethook(), '/lalala')
+
+    def test_set_login_function(self):
+        from machammer.decorators import login
+
+        @login
+        def blaa():
+            import sys
+            import subprocess
+            subprocess.call(['/usr/bin/say', 'Hello ' + sys.argv[1]])
+
+        blaa()
+
+    def test_unset_login(self):
+        hooks.login()
+        with self.assertRaises(Exception):
+            self.assertEquals(self.gethook(), '')
+
+
 if __name__ == '__main__':
     logging.basicConfig(level=logging.DEBUG)
     main()