Added curl and other stuff

6 files changed, 130 insertions, 47 deletions

diff --git a/README.md b/README.md
index 42eb8fb..3e06feb 100644
--- a/README.md
+++ b/README.md
@@ -146,9 +146,9 @@ Check `tests.py` for more usage examples.
 ### FAQ
 
 * Q: Why not use Bash?
-* A: It's true that most of this stuff is just glue to various command line utilities and using Bash might save some keystrokes, but Python is just a much better programming language with an actual standard library. 
+* A: It's true that most of this stuff is just glue to various command line utilities and using Bash might save some keystrokes, but Python is just a much better programming language with an actual standard library.
 * Q: Why not use Munki?
-* A: No reason whatsoever. Munki is great and you should totally use it, if it works for you. I just prefer to read and write code than learn a new XML syntax. For me personally, it was difficult to "start small" with Munki - there's a lot you have to learn to get started. Also, there are plenty of apps out there that don't conform to the standard PKG/app bundle format (like the ArchiCAD example above) and your best bet at tackling those is just plain-old scripting. To paraphrase Einstein - an installation tool might take you from A to B, but scripting can take you anywhere. :-)
+* A: No reason whatsoever. Munki is great and you should totally use it, if it works for you. I just prefer to read and write code than learn a new XML syntax. For me personally, it was difficult to "start small" with Munki - there's a lot you have to learn to get started. Also, there are plenty of apps out there that don't conform to the standard PKG/app bundle format (like the ArchiCAD example above) and your best bet at tackling those is just plain-old scripting. To paraphrase Einstein - an installation tool might take you from A to B, but scripting can take you anywhere. :-) You can think of `machammer` as a tool to create your own version of Munki.
 
 [![Documentation Status](https://readthedocs.org/projects/machammer/badge/?version=latest)](http://machammer.readthedocs.io/en/latest/?badge=latest)
 
diff --git a/machammer/decorators.py b/machammer/decorators.py
index 32ec511..b73b52f 100644
--- a/machammer/decorators.py
+++ b/machammer/decorators.py
@@ -7,23 +7,52 @@ import inspect
 
 import hooks
 
+APP_SUPPORT_DIR = '/Library/Application Support/'
+APP_ID = 'com.github.filipp.machammer'
+
+
+def get_app_support():
+    fp = os.path.join(APP_SUPPORT_DIR, APP_ID)
+
+    if not os.path.exists(fp):
+        os.mkdir(fp, mode=07555)
+
+    return fp
+
+
+def makesource(func):
+    """Build runnable source code from function func."""
+    # the resulting source code
+    r = []
+    # 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)
+    r.insert(0, '# -*- coding: utf-8 -*-')
+    r.insert(0, '#!/usr/bin/env python')
+
+    for l in s:
+        # strip functions indent level
+        r.append(l.replace(indent, ''))
+
+    return '\n'.join(r)
+
+
+def writesource(func, path):
+    """Write source of fun to file at path."""
+    f = open(path, 'w')
+    s = makesource(func)
+    f.write(s)
+    f.close()
+
 
 def login(func):
+    """Install Python function as a LoginHook."""
     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')
+        writesource(func, path)
 
-        f.close()
         # only root should read and execute
         os.chown(path, 0, 0)
         os.chmod(path, stat.S_IXUSR | stat.S_IRUSR)
@@ -36,3 +65,15 @@ def login(func):
         return path
 
     return func_wrapper
+
+
+def agent(func):
+    """Install Python function as a LaunchAgent."""
+    def func_wrapper():
+        path = os.path.join(get_app_support(), 'launchagent.py')
+        writesource(func, path)
+        os.chmod(path, stat.S_IXUSR | stat.S_IRUSR)
+
+        return path
+
+    return func_wrapper
diff --git a/machammer/functions.py b/machammer/functions.py
index 1a7b1d1..522c517 100644
--- a/machammer/functions.py
+++ b/machammer/functions.py
@@ -2,9 +2,10 @@
 
 import os
 import sys
+import logging
 import plistlib
-import subprocess
 import tempfile
+import subprocess
 
 from xml.parsers.expat import ExpatError
 from system_profiler import SystemProfile
@@ -60,12 +61,12 @@ def display_notification(msg, title='', subtitle=''):
 
 def ditto(src, dst):
     """Shortcut for ditto."""
-    subprocess.call(['/usr/bin/ditto', src, dst])
+    call('/usr/bin/ditto', src, dst)
 
 
 def rsync(src, dst, flags='auE'):
     """Shortcut for rsync."""
-    subprocess.call(['/usr/bin/rsync', '-' + flags, src, dst])
+    call('/usr/bin/rsync', '-' + flags, src, dst)
 
 
 def dscl(domain='.', *args):
@@ -76,7 +77,7 @@ def dscl(domain='.', *args):
 def exec_jar(path, user):
     javapath = '/Library/Internet Plug-Ins/JavaAppletPlugin.plugin/Contents/Home/bin/java'
     if not os.path.exists(javapath):
-        raise ValueError('Looks like your machine does not have Java installed')
+        raise Exception('Looks like your machine does not have Java installed')
 
     call('/bin/launchctl', 'asuser', user, javapath, '-jar', path, '-silent')
 
@@ -136,25 +137,25 @@ def mount_image(dmg):
                          stdout=subprocess.PIPE,
                          stdin=subprocess.PIPE,
                          stderr=subprocess.PIPE)
-    r, e = p.communicate(input=b'Q\nY\n')
+    # work around EULA prompt
+    out, err = p.communicate(input=b'Q\nY\n')
+    logging.debug('mount_image got %s' % out)
 
-    if e:
-        raise Exception(e)
+    if err:
+        raise Exception(err)
 
     try:
-        plist = plistlib.readPlistFromString(r)
-    except ExpatError:  # probably a EULA-image, return None instead of breaking
-        return None
-
-    for p in [x.get('mount-point') for x in plist.get('system-entities')]:
-        if p and os.path.exists(p):
-            return p
-
-    raise Exception('Failed to mount %s' % dmg)
+        _, xml = out.split('<?xml version="1.0" encoding="UTF-8"?>')
+        plist = plistlib.readPlistFromString(xml)
+        for p in [x.get('mount-point') for x in plist.get('system-entities')]:
+            if p and os.path.exists(p):
+                return p
+    except ExpatError:
+        raise Exception('Failed to mount %s' % dmg)
 
 
 def mount_and_install(dmg, pkg):
-    """Mountsthe DMG and installs the PKG."""
+    """Mounts the DMG and installs the PKG."""
     p = mount_image(dmg)
     install_pkg(os.path.join(p, pkg))
 
@@ -166,7 +167,15 @@ def install_profile(path):
 
 def install_pkg(pkg, target='/'):
     """Install a package."""
-    subprocess.call(['/usr/sbin/installer', '-pkg', pkg, '-target', target])
+    call('/usr/sbin/installer', '-pkg', pkg, '-target', target)
+
+
+def mount_url(url):
+    """Mount disk image from URL.
+    Return path to mounted volume."""
+    if url.startswith('http'):
+        p = curl(url)
+        return mount_image(p)
 
 
 def mount_afp(url, username, password, mountpoint=None):
@@ -180,11 +189,12 @@ def mount_afp(url, username, password, mountpoint=None):
 
 def umount(path):
     """Unmount path."""
-    subprocess.call(['/sbin/umount', path])
+    call('/sbin/umount', path)
 
 
 def install_su(restart=True):
-    """Install all available Apple software Updates, restart if any 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')
@@ -192,12 +202,9 @@ def install_su(restart=True):
 
 
 def disable_wifi(port='en1'):
-    call('/usr/sbin/networksetup', '-setairportpower', port, 'off')
-    call('/usr/sbin/networksetup', '-setnetworkserviceenabled', 'Wi-Fi', 'off')
-
-
-def log(msg):
-    print('*** %s...' % msg)
+    ns = '/usr/sbin/networksetup'
+    call(ns, '-setairportpower', port, 'off')
+    call(ns, '-setnetworkserviceenabled', 'Wi-Fi', 'off')
 
 
 def install_service(src):
@@ -216,3 +223,14 @@ def clear_xattr(path):
 def create_os_media(src, dst):
     fp = os.path.join(src, 'Contents/Resources/createinstallmedia')
     call(fp, '--volume', dst, '--applicationpath', src, '--nointeraction')
+
+
+def curl(url, *args):
+    """Fetch URL with curl"""
+    dst = tempfile.NamedTemporaryFile(delete=False)
+    call('/usr/bin/curl', url, '-o', dst.name, '--silent', *args)
+    return dst.name
+
+
+def log(msg):
+    print('*** %s...' % msg)
diff --git a/machammer/hooks.py b/machammer/hooks.py
index 40b7986..f975aca 100644
--- a/machammer/hooks.py
+++ b/machammer/hooks.py
@@ -34,10 +34,17 @@ def shutdown(path=None):
 def mount(mountpoint, path=None):
     """Execute path if mountpoint is mounted. Set path to None to disable."""
     label = 'com.github.filipp.machammer.mounthook'
-    defaults.set('/Library/LaunchAgents/%s' % label, 'Label', label)
-    defaults.set('/Library/LaunchAgents/%s' % label, 'Program', path)
-    defaults.set('/Library/LaunchAgents/%s' % label, 'StartOnMount',
-                 '-boolean',
-                 'TRUE')
-    defaults.set('/Library/LaunchAgents/%s' % label, 'WatchPaths', '-array',
-                 mountpoint)
+    plist = '/Library/LaunchDaemons/' + label
+    defaults.set(plist, 'Label', label)
+    defaults.set(plist, 'Program', path)
+    defaults.set(plist, 'StartOnMount', '-boolean', 'TRUE')
+    defaults.set(plist, 'WatchPaths', '-array', mountpoint)
+
+
+def agent(path=None):
+    """Execute path if mountpoint is mounted. Set path to None to disable."""
+    label = 'com.github.filipp.machammer.loginhook'
+    plist = '/Library/LaunchAgents/' + label
+    defaults.set(plist, 'Label', label)
+    defaults.set(plist, 'Program', path)
+    defaults.set(plist, 'RunAtLoad', '-boolean', 'TRUE')
diff --git a/machammer/processes.py b/machammer/process.py
index 71c2936..d445e07 100755
--- a/machammer/processes.py
+++ b/machammer/process.py
@@ -1,14 +1,19 @@
-from .functions import osascript
+# -*- coding: utf-8 -*-
+
+from .functions import osascript, call
 
 
 def pidof(name):
     pass
 
+
 def is_running(name):
     pass
 
+
 def quit(name):
     pass
 
+
 def kill(name):
     pass
diff --git a/tests.py b/tests.py
index 9c9580f..85eee77 100755
--- a/tests.py
+++ b/tests.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
 
+import os
 import time
 import logging
 import subprocess
@@ -106,6 +107,14 @@ class FunctionsTestCase(TestCase):
     def test_sleep(self):
         functions.sleep()
 
+    def test_curl(self):
+        p = functions.curl(os.getenv('MH_URL'))
+        print(p)
+
+    def test_mount_url(self):
+        p = functions.mount_url(os.getenv('MH_URL'))
+        self.assertTrue(os.path.isdir(p))
+
 
 class ScreenSaverTestCase(TestCase):
     def test_set_invalid(self):
@@ -139,6 +148,9 @@ class HooksTestCase(TestCase):
         blaa()
         self.assertEquals(self.gethook(), '/var/root/Library/mh_loginhook.py')
 
+    def test_launchagent(self):
+        pass
+
     def test_unset_login(self):
         hooks.login()
         with self.assertRaises(Exception):