// -*- coding: utf-8-unix -*-
// ---------------------------------------------------------
//  Block Mount Utility
//  Copyright (c) 2013 Decomo
// ---------------------------------------------------------
//  License:
//      2-clause BSD license
//  Reference:
//      https://github.com/nanoant/mountblockd
//      http://stackoverflow.com/questions/3720503/how-can-i-prevent-ejection-of-a-disk-during-an-operation-on-mac-os-x
//      http://superuser.com/questions/336455/mac-lion-fstab-is-deprecated-so-what-replaces-it-to-prevent-a-partition-from-m
//      http://www.geekpage.jp/blog/?id=2011/10/4/3
//      http://developer.apple.com/library/mac/#documentation/Darwin/Reference/DiscArbitrationFramework/
//  Compile:
//      cc BlockMount.c -g -o BlockMount -framework CoreFoundation -framework DiskArbitration
// ---------------------------------------------------------

#include <CoreFoundation/CoreFoundation.h>
#include <DiskArbitration/DiskArbitration.h>
#include <stdio.h>
#include <stdlib.h>

const char **g_blockDeviceList = NULL;
int g_blockDeviceListCount = 0;
int g_verbose = 0;
int g_run = 1;

DADissenterRef diskWillMount(DADiskRef disk, void *context)
{
    const char *devName = DADiskGetBSDName(disk);
    DADissenterRef dissenter = NULL;

    if (devName)
    {
        int block=0;
        for (int i=0; i<g_blockDeviceListCount; i++)
        {
            if (strstr(devName, g_blockDeviceList[i]))
                block=1;
        }
        if (block)
        {
            if (g_verbose)
                fprintf(stderr, "BLOCKED mount `%s'\n", devName);
            dissenter = DADissenterCreate(kCFAllocatorDefault, kDAReturnNotPermitted, NULL);
        }
    }

    return dissenter;
}

void signalHandler(int sig)
{
    switch (sig) {
    case SIGHUP:
        if (g_verbose)
            fprintf(stderr, "received SIGHUP, terminating...\n");
        CFRunLoopStop(CFRunLoopGetCurrent());
        g_run=0;
        break;
    case SIGTERM:
        if (g_verbose)
            fprintf(stderr, "received SIGTERM, terminating...\n");
        CFRunLoopStop(CFRunLoopGetCurrent());
        g_run=0;
        break;
    case SIGINT:
        if (g_verbose)
            fprintf(stderr, "received SIGINT, terminating...\n");
        CFRunLoopStop(CFRunLoopGetCurrent());
        g_run=0;
        break;
    case SIGQUIT:
        if (g_verbose)
            fprintf(stderr, "received SIGQUIT, terminating...\n");
        CFRunLoopStop(CFRunLoopGetCurrent());
        g_run=0;
        break;
    default:
        fprintf(stderr, "unhandled signal (%d) %s\n", sig, strsignal(sig));
        break;
    }
}

int main(int argc, const char *argv[])
{
    int usage=1;
    if (argc >= 2)
    {
        usage=0;
        if (strcmp(argv[1], "-v") == 0)
        {
            g_verbose=1;
            g_blockDeviceList = &argv[2];
            g_blockDeviceListCount = argc-2;
            if (argc == 2)
            {
                usage=1;
            }
        }
        else
        {
            g_blockDeviceList = &argv[1];
            g_blockDeviceListCount = argc-1;
        }
    }

    if (usage)
    {
        fprintf(stderr, "usage: %s [-v] deviceName1 [deviceName2] [deviceName3] [...] \n", argv[0]);
        return EXIT_FAILURE;
    }

    fprintf(stderr, "blocking mount device(s) : ");
    for (int i=0; i<g_blockDeviceListCount; i++)
    {
        fprintf(stderr, "%s ", g_blockDeviceList[i]);
    }
    fprintf(stderr, "\n");

    signal(SIGHUP,  signalHandler);
    signal(SIGTERM, signalHandler);
    signal(SIGINT,  signalHandler);
    signal(SIGQUIT, signalHandler);

    DAApprovalSessionRef session = DAApprovalSessionCreate(kCFAllocatorDefault);
    if (!session)
    {
        fprintf(stderr, "failed to create DAApprovalSession.\n");
        return EXIT_FAILURE;
    }
    else
    {
        DARegisterDiskMountApprovalCallback(session, NULL, diskWillMount, NULL);
        DAApprovalSessionScheduleWithRunLoop(session, CFRunLoopGetCurrent(), kCFRunLoopDefaultMode);

        while (g_run)
        {
            CFRunLoopRun();
        }

        DAApprovalSessionUnscheduleFromRunLoop(session, CFRunLoopGetCurrent(), kCFRunLoopDefaultMode);
        DAUnregisterApprovalCallback(session, diskWillMount, NULL);
        CFRelease(session);
    }

    return 0;
}