/*
 * Copyright (c) 2010 SURFnet bv
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER
 * IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
 * OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN
 * IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/*****************************************************************************
 OSToken.cpp

 The token class; a token is stored in a directory containing several files.
 Each object is stored in a separate file and a token object is present that
 has the token specific attributes
 *****************************************************************************/

#include "config.h"
#include "log.h"
#include "OSAttributes.h"
#include "OSAttribute.h"
#include "ObjectFile.h"
#include "Directory.h"
#include "Generation.h"
#include "UUID.h"
#include "cryptoki.h"
#include "OSToken.h"
#include "OSPathSep.h"
#include <vector>
#include <string>
#include <set>
#include <map>
#include <list>
#include <stdio.h>

// Constructor
OSToken::OSToken(const std::string inTokenPath)
{
	tokenPath = inTokenPath;

	tokenDir = new Directory(tokenPath);
	gen = Generation::create(tokenPath + OS_PATHSEP + "generation", true);
	tokenObject = new ObjectFile(this, tokenPath + OS_PATHSEP + "token.object", tokenPath + OS_PATHSEP + "token.lock");
	tokenMutex = MutexFactory::i()->getMutex();
	valid = (gen != NULL) && (tokenMutex != NULL) && tokenDir->isValid() && tokenObject->valid;

	DEBUG_MSG("Opened token %s", tokenPath.c_str());

	index(true);
}

// Create a new token
/*static*/ OSToken* OSToken::createToken(const std::string basePath, const std::string tokenDir, const ByteString& label, const ByteString& serial)
{
	Directory baseDir(basePath);

	if (!baseDir.isValid())
	{
		ERROR_MSG("Could not create the Directory object");
		return NULL;
	}

	// Create the token directory
	if (!baseDir.mkdir(tokenDir))
	{
		// Error msg is generated by mkdir
		return NULL;
	}

	// Create the token object
	ObjectFile tokenObject(NULL, basePath + OS_PATHSEP + tokenDir + OS_PATHSEP + "token.object", basePath + OS_PATHSEP + tokenDir + OS_PATHSEP + "token.lock", true);

	if (!tokenObject.valid)
	{
		std::string tokenPath = basePath + OS_PATHSEP + tokenDir + OS_PATHSEP + "token.[object|lock]";
		ERROR_MSG("Failed to create the token object: %s", tokenPath.c_str());

		baseDir.rmdir(tokenDir);

		return NULL;
	}

	// Set the initial attributes
	CK_ULONG flags =
		CKF_RNG |
		CKF_LOGIN_REQUIRED | // FIXME: check
		CKF_RESTORE_KEY_NOT_NEEDED |
		CKF_TOKEN_INITIALIZED |
		CKF_SO_PIN_LOCKED |
		CKF_SO_PIN_TO_BE_CHANGED;

	OSAttribute tokenLabel(label);
	OSAttribute tokenSerial(serial);
	OSAttribute tokenFlags(flags);

	if (!tokenObject.setAttribute(CKA_OS_TOKENLABEL, tokenLabel) ||
	    !tokenObject.setAttribute(CKA_OS_TOKENSERIAL, tokenSerial) ||
	    !tokenObject.setAttribute(CKA_OS_TOKENFLAGS, tokenFlags))
	{
		ERROR_MSG("Failed to set the token attributes");

		baseDir.remove(tokenDir + OS_PATHSEP + "token.object");
		baseDir.remove(tokenDir + OS_PATHSEP + "token.lock");
		baseDir.rmdir(tokenDir);

		return NULL;
	}

	DEBUG_MSG("Created new token %s", tokenDir.c_str());

	return new OSToken(basePath + OS_PATHSEP + tokenDir);
}

// Access an existing token
/*static*/ OSToken *OSToken::accessToken(const std::string &basePath, const std::string &tokenDir)
{
	return new OSToken(basePath + OS_PATHSEP + tokenDir);
}

// Destructor
OSToken::~OSToken()
{
	// Clean up
	std::set<OSObject*> cleanUp = allObjects;
	allObjects.clear();

	for (std::set<OSObject*>::iterator i = cleanUp.begin(); i != cleanUp.end(); i++)
	{
		delete *i;
	}

	delete tokenDir;
	if (gen != NULL) delete gen;
	MutexFactory::i()->recycleMutex(tokenMutex);
	delete tokenObject;
}

// Set the SO PIN
bool OSToken::setSOPIN(const ByteString& soPINBlob)
{
	if (!valid) return false;

	OSAttribute soPIN(soPINBlob);

	CK_ULONG flags;

	if (tokenObject->setAttribute(CKA_OS_SOPIN, soPIN) &&
	    getTokenFlags(flags))
	{
		flags &= ~CKF_SO_PIN_COUNT_LOW;
		flags &= ~CKF_SO_PIN_FINAL_TRY;
		flags &= ~CKF_SO_PIN_LOCKED;
		flags &= ~CKF_SO_PIN_TO_BE_CHANGED;

		return setTokenFlags(flags);
	}

	return false;
}

// Get the SO PIN
bool OSToken::getSOPIN(ByteString& soPINBlob)
{
	if (!valid || !tokenObject->isValid())
	{
		return false;
	}

	if (tokenObject->attributeExists(CKA_OS_SOPIN))
	{
		soPINBlob = tokenObject->getAttribute(CKA_OS_SOPIN).getByteStringValue();

		return true;
	}
	else
	{
		return false;
	}
}

// Set the user PIN
bool OSToken::setUserPIN(ByteString userPINBlob)
{
	if (!valid) return false;

	OSAttribute userPIN(userPINBlob);

	CK_ULONG flags;

	if (tokenObject->setAttribute(CKA_OS_USERPIN, userPIN) &&
	    getTokenFlags(flags))
	{
		flags |= CKF_USER_PIN_INITIALIZED;
		flags &= ~CKF_USER_PIN_COUNT_LOW;
		flags &= ~CKF_USER_PIN_FINAL_TRY;
		flags &= ~CKF_USER_PIN_LOCKED;
		flags &= ~CKF_USER_PIN_TO_BE_CHANGED;

		return setTokenFlags(flags);
	}

	return false;
}

// Get the user PIN
bool OSToken::getUserPIN(ByteString& userPINBlob)
{
	if (!valid || !tokenObject->isValid())
	{
		return false;
	}

	if (tokenObject->attributeExists(CKA_OS_USERPIN))
	{
		userPINBlob = tokenObject->getAttribute(CKA_OS_USERPIN).getByteStringValue();

		return true;
	}
	else
	{
		return false;
	}
}

// Retrieve the token label
bool OSToken::getTokenLabel(ByteString& label)
{
	if (!valid || !tokenObject->isValid())
	{
		return false;
	}

	if (tokenObject->attributeExists(CKA_OS_TOKENLABEL))
	{
		label = tokenObject->getAttribute(CKA_OS_TOKENLABEL).getByteStringValue();

		return true;
	}
	else
	{
		return false;
	}
}

// Retrieve the token serial
bool OSToken::getTokenSerial(ByteString& serial)
{
	if (!valid || !tokenObject->isValid())
	{
		return false;
	}

	if (tokenObject->attributeExists(CKA_OS_TOKENSERIAL))
	{
		serial = tokenObject->getAttribute(CKA_OS_TOKENSERIAL).getByteStringValue();

		return true;
	}
	else
	{
		return false;
	}
}

// Get the token flags
bool OSToken::getTokenFlags(CK_ULONG& flags)
{
	if (!valid || !tokenObject->isValid())
	{
		return false;
	}

	if (tokenObject->attributeExists(CKA_OS_TOKENFLAGS))
	{
		flags = tokenObject->getAttribute(CKA_OS_TOKENFLAGS).getUnsignedLongValue();

		// Check if the user PIN is initialised
		if (tokenObject->attributeExists(CKA_OS_USERPIN))
		{
			flags |= CKF_USER_PIN_INITIALIZED;
		}

		return true;
	}
	else
	{
		return false;
	}
}

// Set the token flags
bool OSToken::setTokenFlags(const CK_ULONG flags)
{
	if (!valid) return false;

	OSAttribute tokenFlags(flags);

	return tokenObject->setAttribute(CKA_OS_TOKENFLAGS, tokenFlags);
}

// Retrieve objects
std::set<OSObject*> OSToken::getObjects()
{
	index();

	// Make sure that no other thread is in the process of changing
	// the object list when we return it
	MutexLocker lock(tokenMutex);

	return objects;
}

void OSToken::getObjects(std::set<OSObject*> &inObjects)
{
	index();

	// Make sure that no other thread is in the process of changing
	// the object list when we return it
	MutexLocker lock(tokenMutex);

	inObjects.insert(objects.begin(),objects.end());
}

// Create a new object
OSObject* OSToken::createObject()
{
	if (!valid) return NULL;

	// Generate a name for the object
	std::string objectUUID = UUID::newUUID();
	std::string objectPath = tokenPath + OS_PATHSEP + objectUUID + ".object";
	std::string lockPath = tokenPath + OS_PATHSEP + objectUUID + ".lock";

	// Create the new object file
	ObjectFile* newObject = new ObjectFile(this, objectPath, lockPath, true);

	if (!newObject->valid)
	{
		ERROR_MSG("Failed to create new object %s", objectPath.c_str());

		delete newObject;

		return NULL;
	}

	// Now add it to the set of objects
	MutexLocker lock(tokenMutex);

	objects.insert(newObject);
	allObjects.insert(newObject);
	currentFiles.insert(newObject->getFilename());

	DEBUG_MSG("(0x%08X) Created new object %s (0x%08X)", this, objectPath.c_str(), newObject);

	gen->update();

	gen->commit();

	return newObject;
}

// Delete an object
bool OSToken::deleteObject(OSObject* object)
{
	if (!valid) return false;

	if (objects.find(object) == objects.end())
	{
		ERROR_MSG("Cannot delete non-existent object 0x%08X", object);

		return false;
	}

	MutexLocker lock(tokenMutex);

	ObjectFile* fileObject = dynamic_cast<ObjectFile*>(object);
	if (fileObject == NULL)
	{
		ERROR_MSG("Object type not compatible with this token class 0x%08X", object);

		return false;
	}

	// Invalidate the object instance
	fileObject->invalidate();

	// Retrieve the filename of the object
	std::string objectFilename = fileObject->getFilename();

	// Attempt to delete the file
	if (!tokenDir->remove(objectFilename))
	{
		ERROR_MSG("Failed to delete object file %s", objectFilename.c_str());

		return false;
	}

	// Retrieve the filename of the lock
	std::string lockFilename = fileObject->getLockname();

	// Attempt to delete the lock
	if (!tokenDir->remove(lockFilename))
	{
		ERROR_MSG("Failed to delete lock file %s", lockFilename.c_str());

		return false;
	}

	objects.erase(object);

	DEBUG_MSG("Deleted object %s", objectFilename.c_str());

	gen->update();

	gen->commit();

	return true;
}

// Checks if the token is consistent
bool OSToken::isValid()
{
	return valid;
}

// Invalidate the token (for instance if it is deleted)
void OSToken::invalidate()
{
	valid = false;
}

// Delete the token
bool OSToken::clearToken()
{
	MutexLocker lock(tokenMutex);

	// Invalidate the token
	invalidate();

	// First, clear out all objects
	objects.clear();

	// Now, delete all files in the token directory
	if (!tokenDir->refresh())
	{
		return false;
	}

	std::vector<std::string> tokenFiles = tokenDir->getFiles();

	for (std::vector<std::string>::iterator i = tokenFiles.begin(); i != tokenFiles.end(); i++)
	{
		if (!tokenDir->remove(*i))
		{
			ERROR_MSG("Failed to remove %s from token directory %s", i->c_str(), tokenPath.c_str());

			return false;
		}
	}

	// Now remove the token directory
	if (!tokenDir->rmdir(""))
	{
		ERROR_MSG("Failed to remove the token directory %s", tokenPath.c_str());

		return false;
	}

	DEBUG_MSG("Token instance %s was succesfully cleared", tokenPath.c_str());

	return true;
}

// Reset the token
bool OSToken::resetToken(const ByteString& label)
{
	CK_ULONG flags;

	if (!getTokenFlags(flags))
	{
		ERROR_MSG("Failed to get the token attributes");

		return false;
	}

	// Clean up
	std::set<OSObject*> cleanUp = getObjects();

	MutexLocker lock(tokenMutex);

	for (std::set<OSObject*>::iterator i = cleanUp.begin(); i != cleanUp.end(); i++)
	{
		ObjectFile* fileObject = dynamic_cast<ObjectFile*>(*i);
		if (fileObject == NULL)
		{
			ERROR_MSG("Object type not compatible with this token class 0x%08X", *i);

			return false;
		}

		// Invalidate the object instance
		fileObject->invalidate();

		// Retrieve the filename of the object
		std::string objectFilename = fileObject->getFilename();

		// Attempt to delete the file
		if (!tokenDir->remove(objectFilename))
		{
			ERROR_MSG("Failed to delete object file %s", objectFilename.c_str());

			return false;
		}

		// Retrieve the filename of the lock
		std::string lockFilename = fileObject->getLockname();

		// Attempt to delete the lock
		if (!tokenDir->remove(lockFilename))
		{
			ERROR_MSG("Failed to delete lock file %s", lockFilename.c_str());

			return false;
		}

		objects.erase(*i);

		DEBUG_MSG("Deleted object %s", objectFilename.c_str());
	}

	// The user PIN has been removed
	flags &= ~CKF_USER_PIN_INITIALIZED;
	flags &= ~CKF_USER_PIN_COUNT_LOW;
	flags &= ~CKF_USER_PIN_FINAL_TRY;
	flags &= ~CKF_USER_PIN_LOCKED;
	flags &= ~CKF_USER_PIN_TO_BE_CHANGED;

	// Set new token attributes
	OSAttribute tokenLabel(label);
	OSAttribute tokenFlags(flags);

	if (!tokenObject->setAttribute(CKA_OS_TOKENLABEL, tokenLabel) ||
	    !tokenObject->setAttribute(CKA_OS_TOKENFLAGS, tokenFlags))
	{
		ERROR_MSG("Failed to set the token attributes");

		return false;
	}

	if (tokenObject->attributeExists(CKA_OS_USERPIN) &&
	    !tokenObject->deleteAttribute(CKA_OS_USERPIN))
	{
		ERROR_MSG("Failed to remove USERPIN");

		return false;
	}

	DEBUG_MSG("Token instance %s was succesfully reset", tokenPath.c_str());

	gen->update();
	gen->commit();

	return true;
}

// Index the token
bool OSToken::index(bool isFirstTime /* = false */)
{
	// No access to object mutable fields before
	MutexLocker lock(tokenMutex);

	// Check if re-indexing is required
	if (!isFirstTime && (!valid || !gen->wasUpdated()))
	{
		DEBUG_MSG("No re-indexing is required");

		return true;
	}

	// Check the integrity
	if (!tokenDir->refresh() || !tokenObject->valid)
	{
		ERROR_MSG("Token integrity check failed");

		valid = false;

		return false;
	}

	DEBUG_MSG("Token %s has changed", tokenPath.c_str());

	// Retrieve the directory listing
	std::vector<std::string> tokenFiles = tokenDir->getFiles();

	// Filter out the objects
	std::set<std::string> newSet;

	for (std::vector<std::string>::iterator i = tokenFiles.begin(); i != tokenFiles.end(); i++)
	{
		if ((i->size() > 7) &&
		    (!(i->substr(i->size() - 7).compare(".object"))) &&
		    (i->compare("token.object")))
		{
			newSet.insert(*i);
		}
		else
		{
			DEBUG_MSG("Ignored file %s", i->c_str());
		}
	}

	// Compute the changes compared to the last list of files
	std::set<std::string> addedFiles;
	std::set<std::string> removedFiles;

	if (!isFirstTime)
	{
		// First compute which files were added
		for (std::set<std::string>::iterator i = newSet.begin(); i != newSet.end(); i++)
		{
			if (currentFiles.find(*i) == currentFiles.end())
			{
				addedFiles.insert(*i);
			}
		}

		// Now compute which files were removed
		for (std::set<std::string>::iterator i = currentFiles.begin(); i != currentFiles.end(); i++)
		{
			if (newSet.find(*i) == newSet.end())
			{
				removedFiles.insert(*i);
			}
		}
	}
	else
	{
		addedFiles = newSet;
	}

	currentFiles = newSet;

	DEBUG_MSG("%d objects were added and %d objects were removed", addedFiles.size(), removedFiles.size());
	DEBUG_MSG("Current directory set contains %d objects", currentFiles.size());

	// Now update the set of objects

	// Add new objects
	for (std::set<std::string>::iterator i = addedFiles.begin(); i != addedFiles.end(); i++)
	{
		if ((i->find_last_of('.') == std::string::npos) ||
		    (i->substr(i->find_last_of('.')) != ".object"))
		{
			continue;
		}

		std::string lockName(*i);
		lockName.replace(lockName.find_last_of('.'), std::string::npos, ".lock");

		// Create a new token object for the added file
		ObjectFile* newObject = new ObjectFile(this, tokenPath + OS_PATHSEP + *i, tokenPath + OS_PATHSEP + lockName);

		// Add the object, even invalid ones.
		// This is so the we can read the attributes once
		// the other process has finished writing to disc.
		DEBUG_MSG("(0x%08X) New object %s (0x%08X) added", this, newObject->getFilename().c_str(), newObject);
		objects.insert(newObject);
		allObjects.insert(newObject);
	}

	// Remove deleted objects
	std::set<OSObject*> newObjects;

	for (std::set<OSObject*>::iterator i = objects.begin(); i != objects.end(); i++)
	{
		ObjectFile* fileObject = dynamic_cast<ObjectFile*>((*i));
		if (fileObject == NULL)
		{
			ERROR_MSG("Object type not compatible with this token class 0x%08X", (*i));

			return false;
		}

		DEBUG_MSG("Processing %s (0x%08X)", fileObject->getFilename().c_str(), *i);

		if (removedFiles.find(fileObject->getFilename()) == removedFiles.end())
		{
			DEBUG_MSG("Adding object %s", fileObject->getFilename().c_str());
			// This object gets to stay in the set
			newObjects.insert(*i);
		}
		else
		{
			fileObject->invalidate();
		}
	}

	// Set the new objects
	objects = newObjects;

	DEBUG_MSG("The token now contains %d objects", objects.size());

	return true;
}