/**
 * ntfs.c - Simple functionality for startup, mounting and unmounting of NTFS-based devices.
 *
 * Copyright (c) 2009 Rhys "Shareese" Koedijk
 * Copyright (c) 2006 Michael "Chishm" Chisholm
 *
 * This program/include file is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License as published
 * by the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program/include file 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., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#ifdef HAVE_STDLIB_H
#include <stdlib.h>
#endif
#ifdef HAVE_ERRNO_H
#include <errno.h>
#endif
#ifdef HAVE_STRING_H
#include <string.h>
#endif

#include "ntfs.h"
#include "ntfsinternal.h"
#include "ntfsfile.h"
#include "ntfsdir.h"
#include "gekko_io.h"
#include "cache.h"

// NTFS device driver devoptab
static const devoptab_t devops_ntfs =
{
    NULL, /* Device name */
    sizeof ( ntfs_file_state ),
    ntfs_open_r,
    ntfs_close_r,
    ntfs_write_r,
    ntfs_read_r,
    ntfs_seek_r,
    ntfs_fstat_r,
    ntfs_stat_r,
    ntfs_link_r,
    ntfs_unlink_r,
    ntfs_chdir_r,
    ntfs_rename_r,
    ntfs_mkdir_r,
    sizeof ( ntfs_dir_state ),
    ntfs_diropen_r,
    ntfs_dirreset_r,
    ntfs_dirnext_r,
    ntfs_dirclose_r,
    ntfs_statvfs_r,
    ntfs_ftruncate_r,
    ntfs_fsync_r,
    NULL /* Device data */
};

void ntfsInit ( void )
{
    static bool isInit = false;

    // Initialise ntfs-3g (if not already done so)
    if ( !isInit )
    {
        isInit = true;

        // Set the log handler
#ifdef NTFS_ENABLE_LOG
        ntfs_log_set_handler( ntfs_log_handler_stderr );
#else
        ntfs_log_set_handler( ntfs_log_handler_null );
#endif
        // Set our current local
        ntfs_set_locale();

    }

    return;
}

int ntfsFindPartitions ( const DISC_INTERFACE *interface, sec_t **partitions )
{
    MASTER_BOOT_RECORD mbr;
    PARTITION_RECORD *partition = NULL;
    sec_t partition_starts[NTFS_MAX_PARTITIONS] = {0};
    int partition_count = 0;
    sec_t part_lba = 0;
    int i;

    union
    {
        u8 buffer[512];
        MASTER_BOOT_RECORD mbr;
        EXTENDED_BOOT_RECORD ebr;
        NTFS_BOOT_SECTOR boot;
    } sector;

    // Sanity check
    if ( !interface )
    {
        errno = EINVAL;
        return -1;
    }
    if ( !partitions )
        return 0;

    // Initialise ntfs-3g
    ntfsInit();

    // Start the device and check that it is inserted
    if ( !interface->startup() )
    {
        errno = EIO;
        return -1;
    }
    if ( !interface->isInserted() )
    {
        return 0;
    }

    // Read the first sector on the device
    if ( !interface->readSectors( 0, 1, &sector.buffer ) )
    {
        errno = EIO;
        return -1;
    }

    // If this is the devices master boot record
    if ( sector.mbr.signature == MBR_SIGNATURE )
    {
        memcpy( &mbr, &sector, sizeof( MASTER_BOOT_RECORD ) );
        ntfs_log_debug( "Valid Master Boot Record found\n" );

        // Search the partition table for all NTFS partitions (max. 4 primary partitions)
        for ( i = 0; i < 4; i++ )
        {
            partition = &mbr.partitions[i];
            part_lba = le32_to_cpu( mbr.partitions[i].lba_start );

            ntfs_log_debug( "Partition %i: %s, sector %d, type 0x%x\n", i + 1,
                            partition->status == PARTITION_STATUS_BOOTABLE ? "bootable (active)" : "non-bootable",
                            part_lba, partition->type );

            // Figure out what type of partition this is
            switch ( partition->type )
            {

                    // Ignore empty partitions
                case PARTITION_TYPE_EMPTY:
                    continue;

                    // NTFS partition
                case PARTITION_TYPE_NTFS:
                    {
                        ntfs_log_debug( "Partition %i: Claims to be NTFS\n", i + 1 );

                        // Read and validate the NTFS partition
                        if ( interface->readSectors( part_lba, 1, &sector ) )
                        {
                            if ( sector.boot.oem_id == NTFS_OEM_ID )
                            {
                                ntfs_log_debug( "Partition %i: Valid NTFS boot sector found\n", i + 1 );
                                if ( partition_count < NTFS_MAX_PARTITIONS )
                                {
                                    partition_starts[partition_count] = part_lba;
                                    partition_count++;
                                }
                            }
                            else
                            {
                                ntfs_log_debug( "Partition %i: Invalid NTFS boot sector, not actually NTFS\n", i + 1 );
                            }
                        }

                        break;

                    }

                    // DOS 3.3+ or Windows 95 extended partition
                case PARTITION_TYPE_DOS33_EXTENDED:
                case PARTITION_TYPE_WIN95_EXTENDED:
                    {
                        ntfs_log_debug( "Partition %i: Claims to be Extended\n", i + 1 );

                        // Walk the extended partition chain, finding all NTFS partitions within it
                        sec_t ebr_lba = part_lba;
                        sec_t next_erb_lba = 0;
                        do
                        {

                            // Read and validate the extended boot record
                            if ( interface->readSectors( ebr_lba + next_erb_lba, 1, &sector ) )
                            {
                                if ( sector.ebr.signature == EBR_SIGNATURE )
                                {
                                    ntfs_log_debug( "Logical Partition @ %d: type 0x%x\n", ebr_lba + next_erb_lba,
                                                    sector.ebr.partition.status == PARTITION_STATUS_BOOTABLE ? "bootable (active)" : "non-bootable",
                                                    sector.ebr.partition.type );

                                    // Get the start sector of the current partition
                                    // and the next extended boot record in the chain
                                    part_lba = ebr_lba + next_erb_lba + le32_to_cpu( sector.ebr.partition.lba_start );
                                    next_erb_lba = le32_to_cpu( sector.ebr.next_ebr.lba_start );

                                    // Check if this partition has a valid NTFS boot record
                                    if ( interface->readSectors( part_lba, 1, &sector ) )
                                    {
                                        if ( sector.boot.oem_id == NTFS_OEM_ID )
                                        {
                                            ntfs_log_debug( "Logical Partition @ %d: Valid NTFS boot sector found\n", part_lba );
                                            if ( sector.ebr.partition.type != PARTITION_TYPE_NTFS )
                                            {
                                                ntfs_log_warning( "Logical Partition @ %d: Is NTFS but type is 0x%x; 0x%x was expected\n", part_lba, sector.ebr.partition.type, PARTITION_TYPE_NTFS );
                                            }
                                            if ( partition_count < NTFS_MAX_PARTITIONS )
                                            {
                                                partition_starts[partition_count] = part_lba;
                                                partition_count++;
                                            }
                                        }
                                    }

                                }
                                else
                                {
                                    next_erb_lba = 0;
                                }
                            }

                        }
                        while ( next_erb_lba );

                        break;

                    }

                    // Unknown or unsupported partition type
                default:
                    {

                        // Check if this partition has a valid NTFS boot record anyway,
                        // it might be misrepresented due to a lazy partition editor
                        if ( interface->readSectors( part_lba, 1, &sector ) )
                        {
                            if ( sector.boot.oem_id == NTFS_OEM_ID )
                            {
                                ntfs_log_debug( "Partition %i: Valid NTFS boot sector found\n", i + 1 );
                                if ( partition->type != PARTITION_TYPE_NTFS )
                                {
                                    ntfs_log_warning( "Partition %i: Is NTFS but type is 0x%x; 0x%x was expected\n", i + 1, partition->type, PARTITION_TYPE_NTFS );
                                }
                                if ( partition_count < NTFS_MAX_PARTITIONS )
                                {
                                    partition_starts[partition_count] = part_lba;
                                    partition_count++;
                                }
                            }
                        }

                        break;

                    }

            }

        }

        // Else it is assumed this device has no master boot record
    }
    else
    {
        ntfs_log_debug( "No Master Boot Record was found!\n" );

        // As a last-ditched effort, search the first 64 sectors of the device for stray NTFS partitions
        for ( i = 0; i < 64; i++ )
        {
            if ( interface->readSectors( i, 1, &sector ) )
            {
                if ( sector.boot.oem_id == NTFS_OEM_ID )
                {
                    ntfs_log_debug( "Valid NTFS boot sector found at sector %d!\n", i );
                    if ( partition_count < NTFS_MAX_PARTITIONS )
                    {
                        partition_starts[partition_count] = i;
                        partition_count++;
                    }
                }
            }
        }

    }

    // Shutdown the device
    /*interface->shutdown();*/

    // Return the found partitions (if any)
    if ( partition_count > 0 )
    {
        *partitions = ( sec_t* )ntfs_alloc( sizeof( sec_t ) * partition_count );
        if ( *partitions )
        {
            memcpy( *partitions, &partition_starts, sizeof( sec_t ) * partition_count );
            return partition_count;
        }
    }

    return 0;
}

int ntfsMountAll ( ntfs_md **mounts, u32 flags )
{
    const INTERFACE_ID *discs = ntfsGetDiscInterfaces();
    const INTERFACE_ID *disc = NULL;
    ntfs_md mount_points[NTFS_MAX_MOUNTS];
    sec_t *partitions = NULL;
    int mount_count = 0;
    int partition_count = 0;
    char name[128];
    int i, j, k;

    // Initialise ntfs-3g
    ntfsInit();

    // Find and mount all NTFS partitions on all known devices
    for ( i = 0; discs[i].name != NULL && discs[i].interface != NULL; i++ )
    {
        disc = &discs[i];
        partition_count = ntfsFindPartitions( disc->interface, &partitions );
        if ( partition_count > 0 && partitions )
        {
            for ( j = 0, k = 0; j < partition_count; j++ )
            {

                // Find the next unused mount name
                do
                {
                    sprintf( name, "%s%i", NTFS_MOUNT_PREFIX, k++ );
                    if ( k >= NTFS_MAX_MOUNTS )
                    {
                        ntfs_free( partitions );
                        errno = EADDRNOTAVAIL;
                        return -1;
                    }
                }
                while ( ntfsGetDevice( name, false ) );

                // Mount the partition
                if ( mount_count < NTFS_MAX_MOUNTS )
                {
                    if ( ntfsMount( name, disc->interface, partitions[j], CACHE_DEFAULT_PAGE_SIZE, CACHE_DEFAULT_PAGE_COUNT, flags ) )
                    {
                        strcpy( mount_points[mount_count].name, name );
                        mount_points[mount_count].interface = disc->interface;
                        mount_points[mount_count].startSector = partitions[j];
                        mount_count++;
                    }
                }

            }
            ntfs_free( partitions );
        }
    }

    // Return the mounts (if any)
    if ( mount_count > 0 && mounts )
    {
        *mounts = ( ntfs_md* )ntfs_alloc( sizeof( ntfs_md ) * mount_count );
        if ( *mounts )
        {
            memcpy( *mounts, &mount_points, sizeof( ntfs_md ) * mount_count );
            return mount_count;
        }
    }

    return 0;
}

int ntfsMountDevice ( const DISC_INTERFACE *interface, ntfs_md **mounts, u32 flags )
{
    const INTERFACE_ID *discs = ntfsGetDiscInterfaces();
    const INTERFACE_ID *disc = NULL;
    ntfs_md mount_points[NTFS_MAX_MOUNTS];
    sec_t *partitions = NULL;
    int mount_count = 0;
    int partition_count = 0;
    char name[128];
    int i, j, k;

    // Sanity check
    if ( !interface )
    {
        errno = EINVAL;
        return -1;
    }

    // Initialise ntfs-3g
    ntfsInit();

    // Find the specified device then find and mount all NTFS partitions on it
    for ( i = 0; discs[i].name != NULL && discs[i].interface != NULL; i++ )
    {
        if ( discs[i].interface == interface )
        {
            disc = &discs[i];
            partition_count = ntfsFindPartitions( disc->interface, &partitions );
            if ( partition_count > 0 && partitions )
            {
                for ( j = 0, k = 0; j < partition_count; j++ )
                {

                    // Find the next unused mount name
                    do
                    {
                        sprintf( name, "%s%i", NTFS_MOUNT_PREFIX, k++ );
                        if ( k >= NTFS_MAX_MOUNTS )
                        {
                            ntfs_free( partitions );
                            errno = EADDRNOTAVAIL;
                            return -1;
                        }
                    }
                    while ( ntfsGetDevice( name, false ) );

                    // Mount the partition
                    if ( mount_count < NTFS_MAX_MOUNTS )
                    {
                        if ( ntfsMount( name, disc->interface, partitions[j], CACHE_DEFAULT_PAGE_SIZE, CACHE_DEFAULT_PAGE_COUNT, flags ) )
                        {
                            strcpy( mount_points[mount_count].name, name );
                            mount_points[mount_count].interface = disc->interface;
                            mount_points[mount_count].startSector = partitions[j];
                            mount_count++;
                        }
                    }

                }
                ntfs_free( partitions );
            }
            break;
        }
    }

    // If we couldn't find the device then return with error status
    if ( !disc )
    {
        errno = ENODEV;
        return -1;
    }

    // Return the mounts (if any)
    if ( mount_count > 0 && mounts )
    {
        *mounts = ( ntfs_md* )ntfs_alloc( sizeof( ntfs_md ) * mount_count );
        if ( *mounts )
        {
            memcpy( *mounts, &mount_points, sizeof( ntfs_md ) * mount_count );
            return mount_count;
        }
    }

    return 0;
}

bool ntfsMount ( const char *name, const DISC_INTERFACE *interface, sec_t startSector, u32 cachePageCount, u32 cachePageSize, u32 flags )
{
    ntfs_vd *vd = NULL;
    gekko_fd *fd = NULL;

    // Sanity check
    if ( !name || !interface )
    {
        errno = EINVAL;
        return false;
    }

    // Initialise ntfs-3g
    ntfsInit();

    // Check that the requested mount name is free
    if ( ntfsGetDevice( name, false ) )
    {
        errno = EADDRINUSE;
        return false;
    }

    // Check that we can at least read from this device
    if ( !( interface->features & FEATURE_MEDIUM_CANREAD ) )
    {
        errno = EPERM;
        return false;
    }

    // Allocate the volume descriptor
    vd = ( ntfs_vd* )ntfs_alloc( sizeof( ntfs_vd ) );
    if ( !vd )
    {
        errno = ENOMEM;
        return false;
    }

    // Setup the volume descriptor
    vd->id = interface->ioType;
    vd->flags = 0;
    vd->uid = 0;
    vd->gid = 0;
    vd->fmask = 0;
    vd->dmask = 0;
    vd->atime = ( ( flags & NTFS_UPDATE_ACCESS_TIMES ) ? ATIME_ENABLED : ATIME_DISABLED );
    vd->showHiddenFiles = ( flags & NTFS_SHOW_HIDDEN_FILES );
    vd->showSystemFiles = ( flags & NTFS_SHOW_SYSTEM_FILES );

    // Allocate the device driver descriptor
    fd = ( gekko_fd* )ntfs_alloc( sizeof( gekko_fd ) );
    if ( !fd )
    {
        ntfs_free( vd );
        errno = ENOMEM;
        return false;
    }

    // Setup the device driver descriptor
    fd->interface = interface;
    fd->startSector = startSector;
    fd->sectorSize = 0;
    fd->sectorCount = 0;
    fd->cachePageCount = cachePageCount;
    fd->cachePageSize = cachePageSize;

    // Allocate the device driver
    vd->dev = ntfs_device_alloc( name, 0, &ntfs_device_gekko_io_ops, fd );
    if ( !vd->dev )
    {
        ntfs_free( fd );
        ntfs_free( vd );
        return false;
    }

    // Build the mount flags
    if ( flags & NTFS_READ_ONLY )
        vd->flags |= MS_RDONLY;
    else
    {
        if ( !( interface->features & FEATURE_MEDIUM_CANWRITE ) )
            vd->flags |= MS_RDONLY;
        if ( ( interface->features & FEATURE_MEDIUM_CANREAD ) && ( interface->features & FEATURE_MEDIUM_CANWRITE ) )
            vd->flags |= MS_EXCLUSIVE;
    }
    if ( flags & NTFS_RECOVER )
        vd->flags |= MS_RECOVER;
    if ( flags & NTFS_IGNORE_HIBERFILE )
        vd->flags |= MS_IGNORE_HIBERFILE;

    if ( vd->flags & MS_RDONLY )
        ntfs_log_debug( "Mounting \"%s\" as read-only\n", name );

    // Mount the device
    vd->vol = ntfs_device_mount( vd->dev, vd->flags );
    if ( !vd->vol )
    {
        switch ( ntfs_volume_error( errno ) )
        {
            case NTFS_VOLUME_NOT_NTFS: errno = EINVALPART; break;
            case NTFS_VOLUME_CORRUPT: errno = EINVALPART; break;
            case NTFS_VOLUME_HIBERNATED: errno = EHIBERNATED; break;
            case NTFS_VOLUME_UNCLEAN_UNMOUNT: errno = EDIRTY; break;
            default: errno = EINVAL; break;
        }
        ntfs_device_free( vd->dev );
        ntfs_free( vd );
        return false;
    }

    // Initialise the volume descriptor
    if ( ntfsInitVolume( vd ) )
    {
        ntfs_umount( vd->vol, true );
        ntfs_free( vd );
        return false;
    }

    // Add the device to the devoptab table
    if ( ntfsAddDevice( name, vd ) )
    {
        ntfsDeinitVolume( vd );
        ntfs_umount( vd->vol, true );
        ntfs_free( vd );
        return false;
    }

    return true;
}

void ntfsUnmount ( const char *name, bool force )
{
    ntfs_vd *vd = NULL;

    // Get the devices volume descriptor
    vd = ntfsGetVolume( name );
    if ( !vd )
        return;

    // Remove the device from the devoptab table
    ntfsRemoveDevice( name );

    // Deinitialise the volume descriptor
    ntfsDeinitVolume( vd );

    // Unmount the volume
    ntfs_umount( vd->vol, force );

    // Free the volume descriptor
    ntfs_free( vd );

    return;
}

const char *ntfsGetVolumeName ( const char *name )
{
    ntfs_vd *vd = NULL;
    //ntfs_attr *na = NULL;
    //ntfschar *ulabel = NULL;
    //char *volumeName = NULL;

    // Sanity check
    if ( !name )
    {
        errno = EINVAL;
        return NULL;
    }

    // Get the devices volume descriptor
    vd = ntfsGetVolume( name );
    if ( !vd )
    {
        errno = ENODEV;
        return NULL;
    }
    return vd->vol->vol_name;
    /*

        // If the volume name has already been cached then just use that
        if (vd->name[0])
            return vd->name;

        // Lock
        ntfsLock(vd);

        // Check if the volume name attribute exists
        na = ntfs_attr_open(vd->vol->vol_ni, AT_VOLUME_NAME, NULL, 0);
        if (!na) {
            ntfsUnlock(vd);
            errno = ENOENT;
            return false;
        }

        // Allocate a buffer to store the raw volume name
        ulabel = ntfs_alloc(na->data_size * sizeof(ntfschar));
        if (!ulabel) {
            ntfsUnlock(vd);
            errno = ENOMEM;
            return false;
        }

        // Read the volume name
        if (ntfs_attr_pread(na, 0, na->data_size, ulabel) != na->data_size) {
            ntfs_free(ulabel);
            ntfsUnlock(vd);
            errno = EIO;
            return false;
        }

        // Convert the volume name to the current local
        if (ntfsUnicodeToLocal(ulabel, na->data_size, &volumeName, 0) < 0) {
            errno = EINVAL;
            ntfs_free(ulabel);
            ntfsUnlock(vd);
            return false;
        }

        // If the volume name was read then cache it (for future fetches)
        if (volumeName)
            strcpy(vd->name, volumeName);

        // Close the volume name attribute
        if (na)
            ntfs_attr_close(na);

        // Clean up
        ntfs_free(volumeName);
        ntfs_free(ulabel);

        // Unlock
        ntfsUnlock(vd);

        return vd->name;
    */
}

bool ntfsSetVolumeName ( const char *name, const char *volumeName )
{
    ntfs_vd *vd = NULL;
    ntfs_attr *na = NULL;
    ntfschar *ulabel = NULL;
    int ulabel_len;

    // Sanity check
    if ( !name )
    {
        errno = EINVAL;
        return false;
    }

    // Get the devices volume descriptor
    vd = ntfsGetVolume( name );
    if ( !vd )
    {
        errno = ENODEV;
        return false;
    }

    // Lock
    ntfsLock( vd );

    // Convert the new volume name to unicode
    ulabel_len = ntfsLocalToUnicode( volumeName, &ulabel ) * sizeof( ntfschar );
    if ( ulabel_len < 0 )
    {
        ntfsUnlock( vd );
        errno = EINVAL;
        return false;
    }

    // Check if the volume name attribute exists
    na = ntfs_attr_open( vd->vol->vol_ni, AT_VOLUME_NAME, NULL, 0 );
    if ( na )
    {

        // It does, resize it to match the length of the new volume name
        if ( ntfs_attr_truncate( na, ulabel_len ) )
        {
            ntfs_free( ulabel );
            ntfsUnlock( vd );
            return false;
        }

        // Write the new volume name
        if ( ntfs_attr_pwrite( na, 0, ulabel_len, ulabel ) != ulabel_len )
        {
            ntfs_free( ulabel );
            ntfsUnlock( vd );
            return false;
        }

    }
    else
    {

        // It doesn't, create it now
        if ( ntfs_attr_add( vd->vol->vol_ni, AT_VOLUME_NAME, NULL, 0, ( u8* )ulabel, ulabel_len ) )
        {
            ntfs_free( ulabel );
            ntfsUnlock( vd );
            return false;
        }

    }

    // Reset the volumes name cache (as it has now been changed)
    vd->name[0] = '\0';

    // Close the volume name attribute
    if ( na )
        ntfs_attr_close( na );

    // Sync the volume node
    if ( ntfs_inode_sync( vd->vol->vol_ni ) )
    {
        ntfs_free( ulabel );
        ntfsUnlock( vd );
        return false;
    }

    // Clean up
    ntfs_free( ulabel );

    // Unlock
    ntfsUnlock( vd );

    return true;
}

const devoptab_t *ntfsGetDevOpTab ( void )
{
    return &devops_ntfs;
}