#!/usr/bin/env python

# coding:utf-8

from functools import wraps

import os

import re

import subprocess

import fire

def root_required(method):

    @wraps(method)

    def wrapper(*args, **kwargs):

        if os.geteuid() != 0:

            exit("You need to have root privileges to run this script.\nPlease try again, this time using 'sudo'. Exiting.")

        return method(*args, **kwargs)

    return wrapper

class MountNTFS(object):

    MOUNT_POINT="/Volumes/NTFS"

    def get_output(self, cmd):

        p = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE)

        p.wait()

        output = p.stdout.read()

        return output

    def require_ntfs3g(self):

        if not self.get_output("which ntfs-3g"):

            exit("ntfs-3g not found, brew install ntfs-3g")

    def get_ntfs_disk(self):

        pat = re.compile(r"disk[\w]+")

        output = self.get_output("diskutil list|grep Windows_NTFS")

        fr = pat.findall(output)

        if fr:

            return fr[0]

    def get_ntfs_dev(self):

        pat = re.compile(r"/dev/disk[\w]+")

        output = self.get_output("mount|grep ntfs")

        fr = pat.findall(output)

        if fr:

            return fr[0]

    def is_mounted(self):

        return self.get_output("mount|grep osxfuse")

    @root_required

    def umount(self, device=None):

        if device is None:

            device = self.MOUNT_POINT

        output = self.get_output("umount %s" % device)

    @root_required

    def mount(self):

        MOUNT_POINT = self.MOUNT_POINT

        if self.is_mounted():

            self.umount()

        disk = self.get_ntfs_disk()

        if not disk:

            print "No NTFS Disk Found."

            exit(1)

        dev = self.get_ntfs_dev()

        if dev:

            self.umount(dev)

        else:

            dev = "/dev/%s" % disk

        if not os.path.isdir(MOUNT_POINT):

            os.makedirs(MOUNT_POINT)

        self.require_ntfs3g()

        mount_cmd = "/usr/local/bin/ntfs-3g %s %s -olocal -oallow_other" % (dev, MOUNT_POINT)

        output = self.get_output(mount_cmd)

if __name__ == '__main__':

    fire.Fire(MountNTFS)