/*
 * Copyright (c) 1999 Apple Computer, Inc. All rights reserved.
 *
 * @APPLE_LICENSE_HEADER_START@
 * 
 * "Portions Copyright (c) 1999 Apple Computer, Inc.  All Rights
 * Reserved.  This file contains Original Code and/or Modifications of
 * Original Code as defined in and that are subject to the Apple Public
 * Source License Version 1.0 (the 'License').  You may not use this file
 * except in compliance with the License.  Please obtain a copy of the
 * License at http://www.apple.com/publicsource and read it before using
 * this file.
 * 
 * The Original Code and all software distributed under the License are
 * distributed on an 'AS IS' basis, WITHOUT WARRANTY OF ANY KIND, EITHER
 * EXPRESS OR IMPLIED, AND APPLE HEREBY DISCLAIMS ALL SUCH WARRANTIES,
 * INCLUDING WITHOUT LIMITATION, ANY WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE OR NON-INFRINGEMENT.  Please see the
 * License for the specific language governing rights and limitations
 * under the License."
 * 
 * @APPLE_LICENSE_HEADER_END@
 */
#include <stdio.h>
#ifndef __linux__
#include <stdlib.h>
#include <unistd.h>
#endif
#include <string.h>
#include <errno.h>
#include <fcntl.h>

#include "io.h"
#include "errors.h"
#ifdef __APPLE__
#include <architecture/alignment.h>
#endif

#import "pdisk.h"
#import "partition_map.h"
#import "dump.h"

#define MIN_PARTITION_SIZE 4

typedef int (*funcPtr_t)();

typedef struct {
    unsigned char *	name;
    funcPtr_t		func;
    int			nargs;
    unsigned char *	description;
} cmdline_option_t;


static int block_size();
static int create_partition();
static int delete_partition();
static int disk_is_partitioned();
static int disk_size();
static int get_partition_of_type();
static int get_partition_with_name();
static int initialize();
static int make_bootable();
static int our_dump();
static int partition_display();
static int split_partition();
static int set_automount();
static int set_writable();

cmdline_option_t optionTable[] = {
    { "-blockSize", block_size, 0, "display the block size used by the map" },
    { "-dump", our_dump, 0, "dump the list of partitions" },
    { "-isDiskPartitioned", disk_is_partitioned, 0, "is the disk partitioned" },
    { "-diskSize", disk_size, 0, "prints the size of the disk in megs" },
    { "-partitionSize", partition_display, 0, "get the partition size in blocks" },
    { "-partitionBase", partition_display, 0, "get the partition base in blocks" },
    { "-partitionType", partition_display, 0, "get the partition type" },
    { "-partitionName", partition_display, 0, "get the partition name" },
    { "-partitionEntry", partition_display, 0, "get the partition name, type, base, and size" },
    { "-splitPartition", split_partition, 0, "split an existing partition in two pieces" },
    { "-createPartition", create_partition, 0, "create a new partition" },
    { "-initialize", initialize, 0, "initialize the partition map" },
    { "-deletePartition", delete_partition, 0, "delete a partition" },
    { "-getPartitionOfType", get_partition_of_type, 0, "get partition of specified type" },
    { "-getPartitionWithName", get_partition_with_name, 0, "get partition with specified name" },
    { "-makeBootable", make_bootable, 0, "make a partition bootable" },
    { "-setAutoMount", set_automount, 0, "set or clear the auto-mount bit for HFS volumes" },
    { "-setWritable", set_writable, 0, "set or clear the writable bit in the pmap" },
    { 0, 0, 0 },
};

static void
command_usage()
{
    fprintf(stderr, 
	    "usage:\tpdisk <raw device> <option>\n"
	    "\twhere <raw device> is the live partition ie. /dev/rdisk0\n"
	    "\tand <option> is one of the following:\n");
    
    {
	cmdline_option_t * tbl_p;

	fprintf(stderr, "Option\n");
	for (tbl_p = optionTable; tbl_p->name; tbl_p++) {
	    fprintf(stderr, "%-24s: %s\n", tbl_p->name, tbl_p->description);
	}
    }
    fprintf(stderr, "For more info on the option, type pdisk <raw device> <option>\n");
}

static int
initialize(char * name)
{
    partition_map_header * map;

    map = create_partition_map(name, NULL);
    if (map == NULL)
	return (CMD_FAIL);
    add_partition_to_map("Apple", kMapType,
			 1, (map->media_size <= 128? 2: 63), map);
    write_partition_map(map);
    close_partition_map(map);
    return (CMD_SUCCESS);
}

static int
block_size(char * name)
{
    partition_map_header *map;
    int junk;

    map = open_partition_map(name, &junk, 0);
    if (map == NULL) {
	return (CMD_FAIL);
    }
    printf("%d\n", map->logical_block);
    close_partition_map(map);
    return (CMD_SUCCESS);
}

static int
our_dump(char * name)
{
    partition_map_header *map;
    int junk;

    map = open_partition_map(name, &junk, 0);
    if (map == NULL) {
	return (CMD_FAIL);
    }

    dump_partition_map(map, 1);

    close_partition_map(map);
    return (CMD_SUCCESS);
}

int
disk_is_partitioned(char * name)
{
    partition_map_header *map;
    int junk;

    map = open_partition_map(name, &junk, 0);
    if (map) {
	close_partition_map(map);
	return (CMD_SUCCESS);
    }
    return (CMD_FAIL);
}

int
disk_size(char * name, int argc, char * * argv)
{
    int				junk;
    partition_map_header * 	map;
    Block0 *			p;
    double			d;

    map = open_partition_map(name, &junk, 0);
    if (!map)
	return (CMD_FAIL);
    p = map->misc;
    if (p->sbSig == BLOCK0_SIGNATURE) {
	d = p->sbBlkCount;
	d = d * p->sbBlkSize / (1024 * 1024);
	printf("%lu\n", ((unsigned long)d));
    }
    close_partition_map(map);
    return (CMD_SUCCESS);
}

int
partition_display(char * name, int argc, char * * argv)
{
    partition_map * 		entry;
    int				junk;
    int				j;
    partition_map_header * 	map;
    int				rv = CMD_FAIL;
    char 			which = argv[0][10];
				/* in "-partitionSize", S is 10th char */

    if (argc < 2) {
	fprintf(stderr, "pdisk: %s <partition number>\n", *argv);
	return (CMD_FAIL);
    }

    map = open_partition_map(name, &junk, 0);
    if (!map) {
	fprintf(stderr, "pdisk: disk is not partitioned\n");
	return (CMD_FAIL);
    }
    j = atoi(argv[1]);
    entry = find_entry_by_disk_address(j, map);
    if (!entry) {
	fprintf(stderr, "pdisk: partition %d does not exist\n", j);
    }
    else {
	DPME * p = entry->data;

	rv = CMD_SUCCESS;

	switch (which) {
	  case 'S':
	    printf("%lu\n", p->dpme_pblocks);
	    break;
	  case 'B':
	    printf("%lu\n", p->dpme_pblock_start);
	    break;
	  case 'N':
	    printf("%s\n", p->dpme_name);
	    break;
	  case 'T':
	    printf("%s\n", p->dpme_type);
	    break;
	  case 'E':
	    printf("%s %s %lu %lu\n", p->dpme_name, p->dpme_type, 
		   p->dpme_pblock_start, p->dpme_pblocks);
	    break;
	  default:
	    rv = CMD_FAIL;
	    fprintf(stderr, "pdisk: unexpected program inconsistency\n");
	    break;
	}
    }
    close_partition_map(map);
    return (rv);
}


int
get_partition_of_type(char * name, int argc, char * * argv)
{
    partition_map * 		entry;
    int				junk;
    partition_map_header * 	map;
    int				rv = CMD_FAIL;

    if (argc < 3) {
	fprintf(stderr, "pdisk: %s <type> <instance>\n", *argv);
	return (CMD_FAIL);
    }
    map = open_partition_map(name, &junk, 0);
    if (!map) {
	fprintf(stderr, "pdisk: disk is not partitioned\n");
	return (CMD_FAIL);
    }
    entry = find_entry_of_type(argv[1], atoi(argv[2]), map);
    if (!entry) {
	fprintf(stderr, "pdisk: no partition number %d of type %s\n", 
		atoi(argv[2]), argv[1]);
    }
    else {
	printf("%ld\n", entry->disk_address);
	rv = CMD_SUCCESS;
    }
    close_partition_map(map);
    return (rv);
}

int
get_partition_with_name(char * name, int argc, char * * argv)
{
    partition_map * 		entry;
    int				junk;
    partition_map_header * 	map;
    int				rv = CMD_FAIL;

    if (argc < 3) {
	fprintf(stderr, "pdisk: %s <name> <instance>\n", *argv);
	return (CMD_FAIL);
    }
    map = open_partition_map(name, &junk, 0);
    if (!map) {
	fprintf(stderr, "pdisk: disk is not partitioned\n");
	return (CMD_FAIL);
    }
    entry = find_entry_with_name(argv[1], atoi(argv[2]), map);
    if (!entry) {
	fprintf(stderr, "pdisk: no partition number %d with name %s\n", 
		atoi(argv[2]), argv[1]);
    }
    else {
	printf("%ld\n", entry->disk_address);
	rv = CMD_SUCCESS;
    }
    close_partition_map(map);
    return (rv);
}

int
make_bootable(char * name, int argc, char * * argv)
{
    partition_map * 		entry;
    int				junk;
    partition_map_header * 	map;
    int				part;
    int				rv = CMD_FAIL;
    /* args are: part block bytes load goto */

    if (argc < 6) {
	fprintf(stderr, "pdisk %s <partno> <boot_addr> <boot_bytes>"
		" <load_addr> <goto_addr>\n", argv[0]);
	return (CMD_FAIL);
    }
    map = open_partition_map(name, &junk, 0);
    if (!map) {
	fprintf(stderr, "pdisk: disk is not partitioned\n");
	return (CMD_FAIL);
    }

    part = atoi(argv[1]);
    entry = find_entry_by_disk_address(part, map);
    if (!entry) {
	fprintf(stderr, "pdisk: partition %d does not exist\n", part);
    }
    else if (strcmp(entry->data->dpme_type, kFreeType) == 0
	     || strcmp(entry->data->dpme_type, kMapType) == 0) {
	fprintf(stderr, "pdisk: can't make partition %d bootable, type %s\n",
	       part, entry->data->dpme_type);
    }
    else {
	u32 	boot_block = atol(argv[2]);
	u32 	boot_bytes = atol(argv[3]);
	u32 	load_addr = atol(argv[4]);
	u32 	goto_addr = atol(argv[5]);
	
	entry->data->dpme_load_addr = (u8 *)load_addr;
	entry->data->dpme_goto_addr = (u8 *)goto_addr;
	entry->data->dpme_boot_block = boot_block;
	entry->data->dpme_boot_bytes = boot_bytes;
	strcpy(entry->data->dpme_process_id, BOOT_PARTITION_PROCESS_ID);
	dpme_bootable_set(entry->data, 1); /* mark it bootable */
	write_partition_map(map);
	rv = CMD_SUCCESS;
    }

    return (rv);

}

static int
split_partition(char * name, int argc, char * * argv)
{
    partition_map * 		entry;
    int				junk;
    partition_map_header * 	map;
    u32 			new_size;
    int				part;
    int				rv = CMD_SUCCESS;
    char *			split_name;
    char *			split_type;

    if (argc < 5) {
	fprintf(stderr, "pdisk %s <partno> <1st part size> <2nd part name> <2nd part type>\n",
	       argv[0]);
	return (CMD_FAIL);
    }
    map = open_partition_map(name, &junk, 0);
    if (!map) {
	fprintf(stderr, "pdisk: no valid partitions exist\n");
	return (CMD_FAIL);
    }
    part = atoi(argv[1]);
    new_size = atoi(argv[2]);
    split_name = argv[3];
    split_type = argv[4];
    entry = find_entry_by_disk_address(part, map);
    rv = CMD_FAIL;
    if (!entry) {
	fprintf(stderr, "pdisk: partition %d does not exist\n", part);
    }
    else if (strcmp(entry->data->dpme_type, kFreeType) == 0
	     || strcmp(entry->data->dpme_type, kMapType) == 0) {
	fprintf(stderr, "pdisk: cannot split partition %d because its type is %s\n",
	       part, entry->data->dpme_type);
    }
    else if (!((new_size > 0)
	       && (new_size 
		   <= (entry->data->dpme_pblocks - MIN_PARTITION_SIZE))
	       && (new_size >= MIN_PARTITION_SIZE))) {
	fprintf(stderr, "pdisk: split size of"
	       " partition %d must be between %lu and %lu\n",
	       part, (u32)MIN_PARTITION_SIZE, 
	       (entry->data->dpme_pblocks - (u32)MIN_PARTITION_SIZE));
    }
    else {
	DPME save_dpme;

	save_dpme = *entry->data;
	delete_partition_from_map(entry);
	add_partition_to_map(save_dpme.dpme_name, save_dpme.dpme_type,
			     save_dpme.dpme_pblock_start, new_size, map);
	add_partition_to_map(split_name, split_type, 
			     save_dpme.dpme_pblock_start + new_size,
			     save_dpme.dpme_pblocks - new_size, map);
	write_partition_map(map);

        entry = find_entry_by_base(save_dpme.dpme_pblock_start + new_size, map);
        if (entry) {
            printf("%lu\n", entry->disk_address);
        }
	rv = CMD_SUCCESS;
    }

    close_partition_map(map);
    return (rv);
}


static int
create_partition(char * name, int argc, char * * argv)
{
    partition_map * 		entry;
    int				junk;
    partition_map_header * 	map;
    int				rv = CMD_SUCCESS;
    u32				part_base;
    char *			part_name;
    u32 			part_size;
    char *			part_type;

    if (argc < 5) {
	fprintf(stderr, "%s <name> <type> <base> <size>\n",
	       argv[0]);
	return (CMD_FAIL);
    }
    map = open_partition_map(name, &junk, 0);
    if (!map) {
	fprintf(stderr, "pdisk: disk is not partitioned\n");
	return (CMD_FAIL);
    }
    part_name = argv[1];
    part_type = argv[2];
    part_base = atoi(argv[3]);
    part_size = atoi(argv[4]);
    rv = CMD_FAIL;
    if (add_partition_to_map(part_name, part_type, part_base, part_size, map) 
	== 1) {
	write_partition_map(map);
	entry = find_entry_by_base(part_base, map);
	if (entry) {
	    printf("%lu\n", entry->disk_address);
	    rv = CMD_SUCCESS;
	}
    }
    close_partition_map(map);
    return (rv);
}

static int
delete_partition(char * name, int argc, char * * argv)
{
    partition_map * 		cur;
    int				junk;
    partition_map_header * 	map;
    int				rv = CMD_SUCCESS;
    u32				part_num;

    if (argc < 2) {
        fprintf(stderr, "%s <part>\n",
               argv[0]);
        return (CMD_FAIL);
    }
    map = open_partition_map(name, &junk, 0);
    if (!map) {
        fprintf(stderr, "pdisk: disk is not partitioned\n");
        return (CMD_FAIL);
    }
    part_num = atoi(argv[1]);
    rv = CMD_FAIL;

    cur = find_entry_by_disk_address(part_num, map);
    if (cur == NULL) {
        fprintf(stderr, "No such partition\n");
    } else {
        delete_partition_from_map(cur);
    }
    write_partition_map(map);
    rv = CMD_SUCCESS;   
    
    close_partition_map(map);
    return (rv);
}

int
set_automount(char * name, int argc, char * * argv)
{
    partition_map * 		entry;
    int				junk;
    partition_map_header * 	map;
    int				part;
    int				rv = CMD_FAIL;
    /* args are: part [1|0] */

    if (argc < 3) {
        fprintf(stderr, "pdisk %s <partno> [1|0]\n", argv[0]);
        return (CMD_FAIL);
    }
    map = open_partition_map(name, &junk, 0);
    if (!map) {
        fprintf(stderr, "pdisk: disk is not partitioned\n");
        return (CMD_FAIL);
    }

    part = atoi(argv[1]);
    entry = find_entry_by_disk_address(part, map);
    if (!entry) {
        fprintf(stderr, "pdisk: partition %d does not exist\n", part);
    }
    else if (strcmp(entry->data->dpme_type, kFreeType) == 0
             || strcmp(entry->data->dpme_type, kMapType) == 0) {
        fprintf(stderr, "pdisk: can't set automount on partition %d, type %s\n",
               part, entry->data->dpme_type);
    }
    else {
        u32 	automount_bit = atol(argv[2]);

        if (automount_bit == 0 || automount_bit == 1)
        {
            if (automount_bit)
                entry->data->dpme_flags |= 0x40000000;
            else
                entry->data->dpme_flags &= ~0x40000000;

            write_partition_map(map);
            rv = CMD_SUCCESS;
        }
	else
	{
            fprintf(stderr, "pdisk %s <partno> [1|0]\n", argv[0]);
            return (CMD_FAIL);
	}
    }

    return (rv);

}

int
set_writable(char * name, int argc, char * * argv)
{
    partition_map * 		entry;
    int				junk;
    partition_map_header * 	map;
    int				part;
    int				rv = CMD_FAIL;
    /* args are: part [1|0] */

    if (argc < 3) {
        fprintf(stderr, "pdisk %s <partno> [1|0]\n", argv[0]);
        return (CMD_FAIL);
    }
    map = open_partition_map(name, &junk, 0);
    if (!map) {
        fprintf(stderr, "pdisk: disk is not partitioned\n");
        return (CMD_FAIL);
    }

    part = atoi(argv[1]);
    entry = find_entry_by_disk_address(part, map);
    if (!entry) {
        fprintf(stderr, "pdisk: partition %d does not exist\n", part);
    }
    else if (strcmp(entry->data->dpme_type, kFreeType) == 0
             || strcmp(entry->data->dpme_type, kMapType) == 0) {
        fprintf(stderr, "pdisk: can't set writable bit on partition %d, type %s\n",
               part, entry->data->dpme_type);
    }
    else {
        u32 	writable_bit = atol(argv[2]);

        if (writable_bit == 0 || writable_bit == 1)
        {
            if (writable_bit)
                entry->data->dpme_flags |= 0x00000020;
            else
                entry->data->dpme_flags &= ~0x00000020;

            write_partition_map(map);
            rv = CMD_SUCCESS;
        }
        else
        {
            fprintf(stderr, "pdisk %s <partno> [1|0]\n", argv[0]);
            return (CMD_FAIL);
        }
    }

    return (rv);

}


int
do_command_line(int argc, char * argv[])
{
    unsigned char * 	devName = *argv;
    cmdline_option_t * 	options_p;

    argv++;
    argc--;
    if (!argc) {
        fprintf(stderr, "pdisk: no command specified\n");
	command_usage();
        return (CMD_FAIL);
    }
    for (options_p = &optionTable[0]; options_p->func; options_p++) {
        if (strcmp(options_p->name, *argv) == 0) {
            return ((options_p->func)(devName, argc, argv));
        }
    }
    fprintf(stderr, "pdisk: command %s not recognized\n", *argv);
    command_usage();
    return (CMD_FAIL);
}