#include <unistd.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <sys/ioctl.h>
#include <sys/types.h>
#include <sys/stat.h>
#include "sg_include.h"
#include "sg_err.h"

/* This code is does a SCSI READ CAPACITY command on the given device
   and outputs the result.

*  Copyright (C) 1999 - 2002 D. Gilbert
*  This program 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, or (at your option)
*  any later version.

   This program is program should work on the 2.0, 2.2 and 2.3 series
   of Linux kernels no matter which of those enviroments it was
   compiled and built under.

   Version 0.53 (20021212)

10 byte READ CAPACITY command:
[0x25][  |lu][    ][    ][    ][    ][    ][    ][cntrl] {ignore PMI mode}

*/

#ifndef SG_GET_RESERVED_SIZE
#define SG_GET_RESERVED_SIZE 0x2272
#endif

#ifndef SG_SET_RESERVED_SIZE
#define SG_SET_RESERVED_SIZE 0x2275
#endif

#ifndef SG_GET_VERSION_NUM
#define SG_GET_VERSION_NUM 0x2282
#endif

#ifndef SG_NEXT_CMD_LEN
#define SG_NEXT_CMD_LEN 0x2283
#endif

#ifndef SG_MAX_SENSE
#define SG_MAX_SENSE 16
#endif


#define SG_RT_UNKN (-1)
#define SG_RT_ORIG 0
#define SG_RT_NEW32 1  /* new driver version 2.1.31 + 2.1.32 */
#define SG_RT_NEW34 2  /* new driver version 2.1.34 and after */
        
typedef struct sg_h_n  /* for "forward" compatibility case */
{
    int pack_len;    /* [o] reply_len (ie useless), ignored as input */
    int reply_len;   /* [i] max length of expected reply (inc. sg_header) */
    int pack_id;     /* [io] id number of packet (use ints >= 0) */
    int result;      /* [o] 0==ok, else (+ve) Unix errno (best ignored) */
    unsigned int twelve_byte:1; 
        /* [i] Force 12 byte command length for group 6 & 7 commands  */
    unsigned int target_status:5;   /* [o] scsi status from target */
    unsigned int host_status:8;     /* [o] host status (see "DID" codes) */
    unsigned int driver_status:8;   /* [o] driver status+suggestion */
    unsigned int other_flags:10;    /* unused */
    unsigned char sense_buffer[SG_MAX_SENSE]; /* [o] Output in 3 cases:
           when target_status is CHECK_CONDITION or 
           when target_status is COMMAND_TERMINATED or
           when (driver_status & DRIVER_SENSE) is true. */
} Sg_h_n;      /* This structure is 36 bytes long on i386 */


static int open_scsi_dev_as_sg(char * devname);

#define OFF sizeof(struct sg_header)
#define RCAP_REPLY_LEN 8    
#define RCAP_CMD_LEN 10
#define MY_PACK_ID 1234


static void usage()
{
    printf("Usage: 'sg_readcap <scsi_device>'\n");
}


int main(int argc, char * argv[])
{
    int sg_fd, k, cmd_len, ok;
    unsigned char rcapCmdBlk [RCAP_CMD_LEN] =
                                {0x25, 0, 0, 0, 0, 0, 0, 0, 0, 0};
    unsigned char rcapBuff[OFF + OFF + RCAP_REPLY_LEN]; /* Overkill */
    int rcapInLen = OFF + sizeof(rcapCmdBlk);
    int rcapOutLen = OFF + RCAP_REPLY_LEN;
    unsigned char * buffp = rcapBuff + OFF;
    struct sg_header * sghp = (struct sg_header *)rcapBuff;
    Sg_h_n * n_sghp = (Sg_h_n *)rcapBuff;
    char * file_name = 0;
    int reserved_size;
    int sg_which = SG_RT_UNKN;
    int sg_version = 0;

    for (k = 1; k < argc; ++k) {
        if (*argv[k] == '-') {
            printf("Unrecognized switch: %s\n", argv[k]);
            file_name = 0;
            break;
        }
        else
            file_name = argv[k];
    }
    if (0 == file_name) {
        usage();
        return 1;
    }
    
    sg_fd = open_scsi_dev_as_sg(file_name);
    if (sg_fd < 0) {
        if (-9999 == sg_fd)
            printf("Failed trying to open SCSI device as an sg device\n");
        else
            perror("sg_readcap: open error");
        return 1;
    }

    /* Run time selection code follows */
    if (ioctl(sg_fd, SG_GET_RESERVED_SIZE, &reserved_size) < 0) {
        reserved_size = SG_BIG_BUFF;
        sg_which = SG_RT_ORIG;
    }
    else if (ioctl(sg_fd, SG_GET_VERSION_NUM, &sg_version) < 0)
        sg_which = SG_RT_NEW32;
    else 
        sg_which = SG_RT_NEW34;

    sghp->reply_len = rcapOutLen;
    sghp->pack_id = MY_PACK_ID;
    sghp->twelve_byte = 0;
    sghp->other_flags = 0;     /* some apps assume this is done ?!? */

    switch (sg_which) {
    case SG_RT_ORIG:
        sghp->sense_buffer[0] = 0;
        break;
    case SG_RT_NEW32:
        break;
    case SG_RT_NEW34: /* this is optional, explicitly setting cmd length */
        cmd_len = RCAP_CMD_LEN; 
        if (ioctl(sg_fd, SG_NEXT_CMD_LEN, &cmd_len) < 0) {
            perror("sg_readcap: SG_NEXT_CMD_LEN error");
            close(sg_fd);
            return 1;
        }
        break;
    default:
        printf("Illegal state for sg_which=%d\n", sg_which);
        return 1;
    }
    memcpy(rcapBuff + OFF, rcapCmdBlk, RCAP_CMD_LEN);

    if (write(sg_fd, rcapBuff, rcapInLen) < 0) {
        perror("sg_readcap: write error");
        close(sg_fd);
        return 1;
    }
    
    if (read(sg_fd, rcapBuff, rcapOutLen) < 0) {
        perror("sg_readcap: read error");
        close(sg_fd);
        return 1;
    }

    /* now for the error processing */
    ok = 0;
    switch (sg_which) {
    case SG_RT_ORIG:
        if ((0 == sghp->result) && (0 == sghp->sense_buffer[0]))
            ok = 1;
        else if (sghp->sense_buffer[0])
            sg_print_sense("READ CAPACITY command error", sghp->sense_buffer,
                           SG_MAX_SENSE);
        else /* sghp->result is != 0 */
            printf("READ CAPACITY failed, sghp->result=%d\n", sghp->result);
        break;
    case SG_RT_NEW32:
    case SG_RT_NEW34: 
        switch (sg_err_category(n_sghp->target_status, n_sghp->host_status,
                n_sghp->driver_status, n_sghp->sense_buffer, SG_MAX_SENSE)) {
        case SG_ERR_CAT_CLEAN:
            ok = 1;
            break;
        case SG_ERR_CAT_RECOVERED:
            printf("Recovered error on READ CAPACITY, continuing\n");
            ok = 1;
            break;
        default: /* won't bother decoding other categories */
            sg_chk_n_print("READ CAPACITY command error", 
                           n_sghp->target_status, 
                           n_sghp->host_status, n_sghp->driver_status, 
                           n_sghp->sense_buffer, SG_MAX_SENSE);
            break;
        }
        break;
    default:
        break;
    }

    if (ok) { /* output result if it is available */
        unsigned int last_blk_addr;
        unsigned int block_size;

        last_blk_addr = 0xff & *(buffp + 3);
        last_blk_addr |= (0xff & *(buffp + 2)) << 8;
        last_blk_addr |= (0xff & *(buffp + 1)) << 16;
        last_blk_addr |= (0xff & *buffp) << 24;
        block_size = 0xff & *(buffp + 7);
        block_size |= (0xff & *(buffp + 6)) << 8;
        block_size |= (0xff & *(buffp + 5)) << 16;
        block_size |= (0xff & *(buffp + 4)) << 24;
        printf("Read Capacity results:\n");
        printf("   Last block address = %u (0x%x), Number of blocks = %u\n", 
               last_blk_addr, (int)last_blk_addr, last_blk_addr + 1);
        printf("   Block size = %u bytes\n", block_size);
    }

    close(sg_fd);
    return 0;
}
    

#define MAX_SG_DEVS 26
#define MAX_FILENAME_LEN 128

#define SCAN_ALPHA 0
#define SCAN_NUMERIC 1
#define DEF_SCAN SCAN_NUMERIC

static void make_dev_name(char * fname, int k, int do_numeric)
{
    char buff[MAX_FILENAME_LEN];

    strcpy(fname, "/dev/sg");
    if (do_numeric) {
        sprintf(buff, "%d", k);
        strcat(fname, buff);
    }
    else {
        if (k <= 26) {
            buff[0] = 'a' + (char)k;
            buff[1] = '\0';
            strcat(fname, buff);
        }
        else
            strcat(fname, "xxxx");
    }
}

typedef struct my_scsi_idlun
{
    int mux4;
    int host_unique_id;

} My_scsi_idlun;

static int open_scsi_dev_as_sg(char * devname)
{
    int fd, bus, bbus, k;
    My_scsi_idlun m_idlun, mm_idlun;
    int do_numeric = DEF_SCAN;
    char name[MAX_FILENAME_LEN];
    struct stat a_st;
    int block_dev = 0;

    strcpy(name, devname);
    if ((fd = open(name, O_RDONLY | O_NONBLOCK)) < 0) {
        if (EACCES == errno) {
            if ((fd = open(name, O_RDWR | O_NONBLOCK)) < 0)
                return fd;
        }
    }
    if (fstat(fd, &a_st) < 0) {
        fprintf(stderr, "could do fstat() on fd ??\n");
        close(fd);
        return -9999;
    }
    if (S_ISBLK(a_st.st_mode)) {
        block_dev = 1;
    }
    if (block_dev || (ioctl(fd, SG_GET_TIMEOUT, 0) < 0)) { /* not sg device */
        if (ioctl(fd, SCSI_IOCTL_GET_BUS_NUMBER, &bus) < 0) {
            printf("Need a filename that resolves to a SCSI device\n");
            close(fd);
            return -9999;
        }
        if (ioctl(fd, SCSI_IOCTL_GET_IDLUN, &m_idlun) < 0) {
            printf("Need a filename that resolves to a SCSI device (2)\n");
            close(fd);
            return -9999;
        }
        close(fd);
    
        for (k = 0; k < MAX_SG_DEVS; k++) {
            make_dev_name(name, k, do_numeric);
            if ((fd = open(name, O_RDONLY | O_NONBLOCK)) < 0) {
                if (EACCES == errno) 
                    fd = open(name, O_RDWR | O_NONBLOCK);
                if (fd < 0) {
                    if ((ENOENT == errno) && (0 == k) && 
                        (do_numeric == DEF_SCAN)) {
                        do_numeric = ! DEF_SCAN;
                        make_dev_name(name, k, do_numeric);
                        if ((fd = open(name, O_RDONLY | O_NONBLOCK)) < 0) {
                            if (EACCES == errno) 
                                fd = open(name, O_RDWR | O_NONBLOCK);
                        }
                    }
                    if (fd < 0) {
                        if (EBUSY == errno)
                            continue;  /* step over if O_EXCL already on it */
                        else
                            break;
                    }
                }
            }
            if (ioctl(fd, SCSI_IOCTL_GET_BUS_NUMBER, &bbus) < 0) {
                perror("sg ioctl failed");
                close(fd);
                fd = -9999;
            }
            if (ioctl(fd, SCSI_IOCTL_GET_IDLUN, &mm_idlun) < 0) {
                perror("sg ioctl failed (2)");
                close(fd);
                fd = -9999;
            }
            if ((bus == bbus) && 
                ((m_idlun.mux4 & 0xff) == (mm_idlun.mux4 & 0xff)) &&
                (((m_idlun.mux4 >> 8) & 0xff) == 
                                        ((mm_idlun.mux4 >> 8) & 0xff)) &&
                (((m_idlun.mux4 >> 16) & 0xff) == 
                                        ((mm_idlun.mux4 >> 16) & 0xff))) {
                printf("  >>> Mapping %s to sg device: %s\n", devname, name);
                break;
            }
            else {
                close(fd);
                fd = -9999;
            }
        }
    }
    if (fd >= 0) { /* everything ok, close and re-open read-write */
        close(fd);
        return open(name, O_RDWR);
    }
    else
        return fd;
}