/*
 * Copyright (C) 2009 Chase Douglas
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as published
 * by the Free Software Foundation.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * In addition, as a special exception, the copyright holders give
 * permission to link the code of portions of this program with the
 * OpenSSL library under certain conditions as described in each
 * individual source file, and distribute linked combinations
 * including the two.
 * You must obey the GNU General Public License in all respects
 * for all of the code used other than OpenSSL.  If you modify
 * file(s) with this exception, you may extend this exception to your
 * version of the file(s), but you are not obligated to do so.  If you
 * do not wish to do so, delete this exception statement from your
 * version.
 */

#include <QByteArray>
#include <QHostAddress>

#import <Foundation/Foundation.h>

#include "Connection.h"
#include "MacAuthThread.h"
#include "rinput.h"

static MacAuthThread *macAuthThread = NULL;
static QMutex macAuthThreadLock;

static void dynamicStoreCallback(SCDynamicStoreRef store, CFArrayRef changedKeys, void *info) {
    macAuthThread->consoleUserChanged();
}

static inline void setupRights(AuthorizationRef auth) {
    // Check if right already exists
    if (AuthorizationRightGet("com.rinputd.control", NULL) != errAuthorizationSuccess) {
        // Right doesn't exist, so create it
        NSDictionary *right = [NSDictionary dictionaryWithObjectsAndKeys:
                               @"rule", @"class",
                               @"authenticate-session-owner-or-admin", @"rule",
                               [NSNumber numberWithBool:false], @"shared",
                               [NSNumber numberWithInt:0], @"timeout",
                               nil];
        OSStatus error = AuthorizationRightSet(auth, "com.rinputd.control", (CFDictionaryRef)right, NULL, NULL, NULL);
        if (error != noErr) {
            qCritical("Error: Failed to set authorization right: %d", (int)error);
        }
    }
}

MacAuthThread::MacAuthThread() :
    auth(NULL),
    dynamicStore(NULL),
    storeRunLoopSource(NULL),
    runLoop(NULL) {}

// Initialize internal state of authorization thread
// Objects referenced here are released when thread exits
bool MacAuthThread::setup() {
    // Ensure only one MacAuthThread is initialized
    QMutexLocker locker(&macAuthThreadLock);
    if (macAuthThread) {
        return FALSE;
    }
    
    // Initialize authorization services structures
    OSStatus error = AuthorizationCreate(NULL, NULL, kAuthorizationFlagDefaults, &auth);
    if (error != noErr) {
        qCritical("Error: Failed to create authorization reference: %d", (int)error);
        return FALSE;
    }
    
    setupRights(auth);
    
    rightItem.valueLength = 0;
    rightItem.value = NULL;
    rightItem.flags = 0;
    
    rights.count = 1;
    rights.items = &rightItem;
    
    environmentItems[0].name = kAuthorizationEnvironmentUsername;
    environmentItems[0].flags = 0;
    environmentItems[1].name = kAuthorizationEnvironmentPassword;
    environmentItems[1].flags = 0;
    
    environment.count = 2;
    environment.items = environmentItems;
    
    // Set up callback for notifications of user switches
    SCDynamicStoreContext context;
    context.version = 0;
    context.info = NULL;
    context.retain = NULL;
    context.release = NULL;
    context.copyDescription = NULL;
    
    dynamicStore = SCDynamicStoreCreate(kCFAllocatorDefault, CFSTR("com.rinputd"), dynamicStoreCallback, &context);
    
    NSArray *keys = [[NSArray alloc] initWithObjects:(NSString *)SCDynamicStoreKeyCreateConsoleUser(kCFAllocatorDefault), nil];
    
    if (!SCDynamicStoreSetNotificationKeys(dynamicStore, (CFArrayRef)keys, NULL)) {
        qCritical("Error: Failed to set console user system configuration notification key");
        [keys release];
        CFRelease(dynamicStore);
        CFRelease(auth);
        return FALSE;
    }
    [keys release];
    
    storeRunLoopSource = SCDynamicStoreCreateRunLoopSource(kCFAllocatorDefault, dynamicStore, 0);
    if (!storeRunLoopSource) {
        qCritical("Error: Failed to create run loop source");
        CFRelease(dynamicStore);
        return FALSE;
    }
    
    macAuthThread = this;
    
    return TRUE;
}

void MacAuthThread::run() {
    if (!storeRunLoopSource || !dynamicStore || !auth) {
        qCritical("Error: MacAuthThread started without proper initialization");
        return;
    }
    
    runLoop = CFRunLoopGetCurrent();
    CFRunLoopAddSource(runLoop, storeRunLoopSource, kCFRunLoopCommonModes);
    CFRunLoopRun();
    
    CFRelease(storeRunLoopSource);
    CFRelease(dynamicStore);
    CFRelease(auth);
}

void MacAuthThread::consoleUserChanged() {
    QMutexLocker locker(&mutex);
    
    QMutableVectorIterator<Connection *> connectionsIterator(connections);
    while (connectionsIterator.hasNext()) {
        Connection *connection = connectionsIterator.next();
        
        qWarning("Revoking authorization for client %s", qPrintable(connection->peerAddress().toString()));
        
        rinput_message_t message = { RINPUT_ERROR };
        message.data.error.type = RINPUT_REVOKED_AUTHORIZATION;
        hton_rinput(&message);
        connection->write(QByteArray::fromRawData((const char *)&message, sizeof(message)));
        
        connection->disconnectLater();
        connectionsIterator.remove();
    }
}

bool MacAuthThread::_authorized(Connection &connection, const QByteArray &user, const QByteArray &pass) {
    QMutexLocker locker(&mutex);
    
    // Check if current console user is this process' owner
    uid_t uid;
    CFStringRef consoleUser = SCDynamicStoreCopyConsoleUser(NULL, &uid, NULL);
    CFRelease(consoleUser);
    
    if (uid != getuid()) {
        return FALSE; // Can't allow this process to control someone else's console
    }
    
    // Verify that user has the right to control this console
    environment.items[0].valueLength = user.length();
    environment.items[0].value = const_cast<char *>(user.data());
    environment.items[1].valueLength = pass.length();
    environment.items[1].value = const_cast<char *>(pass.data());
    
    rights.items[0].name = "com.rinputd.control";
    OSStatus error = AuthorizationCopyRights(auth, &rights, &environment, kAuthorizationFlagExtendRights | kAuthorizationFlagDestroyRights, NULL);
    if (error == errAuthorizationSuccess) {
        connections.append(&connection);
        return TRUE;
    }
    else if (error != errAuthorizationDenied && error != errAuthorizationInteractionNotAllowed) {
        qCritical("Error: Authentication processing for right com.rinputd.control.console failed: %d", (int)error);
    }
    
    return FALSE;
}

MacAuthThread::~MacAuthThread() {
    // Stop thread, resources used there will be freed by the thread
    if (runLoop) {
        CFRunLoopStop(runLoop);
    }
    
    QMutexLocker locker(&mutex);
    
    // Close all open connections
    QMutableVectorIterator<Connection *> connectionsIterator(connections);
    while (connectionsIterator.hasNext()) {
        Connection *connection = connectionsIterator.next();
        connection->disconnectLater();
        connectionsIterator.remove();
    }
}

bool MacAuthThread::authorized(Connection &connection, const QByteArray &user, const QByteArray &pass) {
    return macAuthThread->_authorized(connection, user, pass);
}
