From 0b4c5b9cb97d4549f699e11dcad20574b53bd86c Mon Sep 17 00:00:00 2001
From: bigeagle <justin.w.xd@gmail.com>
Date: Tue, 9 Dec 2014 13:27:56 +0800
Subject: [PATCH] add  and  options

---
 examples/tunasync.conf      |  1 +
 tunasync/exec_pre_post.py   | 35 +++++++++++++++++++++++++++++++++++
 tunasync/jobs.py            |  1 +
 tunasync/mirror_config.py   | 10 ++++++++++
 tunasync/mirror_provider.py |  3 ++-
 5 files changed, 49 insertions(+), 1 deletion(-)
 create mode 100644 tunasync/exec_pre_post.py

diff --git a/examples/tunasync.conf b/examples/tunasync.conf
index cf949a1..61d114e 100644
--- a/examples/tunasync.conf
+++ b/examples/tunasync.conf
@@ -31,6 +31,7 @@ provider = "shell"
 command = "sleep 10"
 local_dir = "/mnt/sdb1/mirror/archlinux/current/"
 # log_file = "/dev/null"
+exec_post_sync = "/bin/bash -c 'date --utc \"+%s\" > ${TUNASYNC_WORKING_DIR}/.timestamp'"
 
 [[mirrors]]
 name = "arch2"
diff --git a/tunasync/exec_pre_post.py b/tunasync/exec_pre_post.py
new file mode 100644
index 0000000..892b673
--- /dev/null
+++ b/tunasync/exec_pre_post.py
@@ -0,0 +1,35 @@
+#!/usr/bin/env python2
+# -*- coding:utf-8 -*-
+import os
+import sh
+import shlex
+from .hook import JobHook
+
+
+class CmdExecHook(JobHook):
+    POST_SYNC = "post_sync"
+    PRE_SYNC = "pre_sync"
+
+    def __init__(self, command, exec_at=POST_SYNC):
+        self.command = shlex.split(command)
+        if exec_at == self.POST_SYNC:
+            self.before_job = self._keep_calm
+            self.after_job = self._exec
+        elif exec_at == self.PRE_SYNC:
+            self.before_job = self._exec
+            self.after_job = self._keep_calm
+
+    def _keep_calm(self, ctx={}, **kwargs):
+        pass
+
+    def _exec(self, ctx={}, **kwargs):
+        new_env = os.environ.copy()
+        new_env["TUNASYNC_MIRROR_NAME"] = ctx["mirror_name"]
+        new_env["TUNASYNC_WORKING_DIR"] = ctx["current_dir"]
+
+        _cmd = self.command[0]
+        _args = [] if len(self.command) == 1 else self.command[1:]
+        cmd = sh.Command(_cmd)
+        cmd(*_args, _env=new_env)
+
+# vim: ts=4 sw=4 sts=4 expandtab
diff --git a/tunasync/jobs.py b/tunasync/jobs.py
index 280c4f1..8656fb3 100644
--- a/tunasync/jobs.py
+++ b/tunasync/jobs.py
@@ -40,6 +40,7 @@ def run_job(sema, child_q, manager_q, provider, **settings):
         manager_q.put(("UPDATE", (provider.name, status)))
         ctx = {}   # put context info in it
         ctx['current_dir'] = provider.local_dir
+        ctx['mirror_name'] = provider.name
 
         try:
             for hook in provider.hooks:
diff --git a/tunasync/mirror_config.py b/tunasync/mirror_config.py
index 3015f48..d294c04 100644
--- a/tunasync/mirror_config.py
+++ b/tunasync/mirror_config.py
@@ -5,6 +5,7 @@ from datetime import datetime
 from .mirror_provider import RsyncProvider, ShellProvider
 from .btrfs_snapshot import BtrfsHook
 from .loglimit import LogLimitHook
+from .exec_pre_post import CmdExecHook
 
 
 class MirrorConfig(object):
@@ -126,6 +127,15 @@ class MirrorConfig(object):
             hooks.append(BtrfsHook(service_dir, working_dir, gc_dir))
 
         hooks.append(LogLimitHook())
+
+        if self.exec_pre_sync:
+            hooks.append(
+                CmdExecHook(self.exec_pre_sync, CmdExecHook.PRE_SYNC))
+
+        if self.exec_post_sync:
+            hooks.append(
+                CmdExecHook(self.exec_post_sync, CmdExecHook.POST_SYNC))
+
         return hooks
 
 # vim: ts=4 sw=4 sts=4 expandtab
diff --git a/tunasync/mirror_provider.py b/tunasync/mirror_provider.py
index 38cb843..2291ed5 100644
--- a/tunasync/mirror_provider.py
+++ b/tunasync/mirror_provider.py
@@ -2,6 +2,7 @@
 # -*- coding:utf-8 -*-
 import sh
 import os
+import shlex
 from datetime import datetime
 
 
@@ -110,7 +111,7 @@ class ShellProvider(MirrorProvider):
         super(ShellProvider, self).__init__(name, local_dir, log_dir, log_file,
                                             interval, hooks)
         self.upstream_url = str(upstream_url)
-        self.command = command.split()
+        self.command = shlex.split(command)
 
     def run(self, ctx={}):