import com.jniwrapper.*;
import com.jniwrapper.win32.FunctionName;
import com.jniwrapper.win32.Handle;
import com.jniwrapper.win32.LastErrorException;
import com.jniwrapper.win32.system.Kernel32;

import java.text.MessageFormat;

/**
 * Example demonstrates the working with <a href = "http://msdn.microsoft.com/en-us/library/aa363216%28v=vs.85%29.aspx">DeviceIoControl</a> API function and
 * corresponding <a href = "http://msdn.microsoft.com/en-us/library/aa363987(v=VS.85).aspx">Disk Management Structures</a> to determine
 * a drive type.
 *
 * @author Serge Piletsky
 */
public class DetectingDriveType
{
    static final int FILE_SHARE_READ = 0x00000001;
    static final int FILE_SHARE_WRITE = 0x00000002;
    static final int OPEN_EXISTING = 3;

    static final int IOCTL_STORAGE_BASE = 0x0000002D;
    static final int METHOD_BUFFERED = 0;
    static final int FILE_ANY_ACCESS = 0;

    static final long IOCTL_STORAGE_QUERY_PROPERTY = CTL_CODE(IOCTL_STORAGE_BASE, 0x0500, METHOD_BUFFERED, FILE_ANY_ACCESS);

    static long CTL_CODE(long deviceType, long function, long method, long access) {
        return (deviceType << 16) | (access << 14) | (function << 2) | method;
    }

    /**
     * Wrapper for <a href = "http://msdn.microsoft.com/en-us/library/ff800833%28v=vs.85%29.aspx">STORAGE_BUS_TYPE</a> enumeration.
     */
    public static class STORAGE_BUS_TYPE extends Int32 {
        public static final int BusTypeUnknown = 0x00;
        public static final int BusTypeScsi = 0x01;
        public static final int BusTypeAtapi = 0x02;
        public static final int BusTypeAta = 0x03;
        public static final int BusType1394 = 0x04;
        public static final int BusTypeSsa = 0x05;
        public static final int BusTypeFibre = 0x06;
        public static final int BusTypeUsb = 0x07;
        public static final int BusTypeRAID = 0x08;
        public static final int BusTypeiSCSI = 0x09;
        public static final int BusTypeSas = 0x0A;
        public static final int BusTypeSata = 0x0B;
        public static final int BusTypeSd = 0xC;
        public static final int BusTypeMmc = 0xD;
        public static final int BusTypeVirtual = 0xE;
        public static final int BusTypeFileBackedVirtual = 0xF;
        public static final int BusTypeMax = 0x10;
        public static final int BusTypeMaxReserved = 0x7F;
    }

    /**
     * Wrapper for <a href = "http://msdn.microsoft.com/en-us/library/ff800839%28v=VS.85%29.aspx">STORAGE_PROPERTY_ID</a> enumeration.
     */
    public static class STORAGE_PROPERTY_ID extends Int {
        public static final int StorageDeviceProperty = 0;
        public static final int StorageAdapterProperty = 1;
        public static final int StorageDeviceIdProperty = 2;
        public static final int StorageDeviceUniqueIdProperty = 3;
        public static final int StorageDeviceWriteCacheProperty = 4;
        public static final int StorageMiniportProperty = 5;
        public static final int StorageAccessAlignmentProperty = 6;
        public static final int StorageDeviceSeekPenaltyProperty = 7;
        public static final int StorageDeviceTrimProperty = 8;
    }

    /**
     * Wrapper for <a href = "http://msdn.microsoft.com/en-us/library/ff800841%28v=VS.85%29.aspx">STORAGE_QUERY_TYPE</a> enumeration.
     */
    public static class STORAGE_QUERY_TYPE extends Int {
        public static final int PropertyStandardQuery = 0;
        public static final int PropertyExistsQuery = 1;
        public static final int PropertyMaskQuery = 2;
        public static final int PropertyQueryMaxDefined = 3;
    }

    /**
     * Wrapper for <a href = "http://msdn.microsoft.com/en-us/library/ff800840%28v=VS.85%29.aspx">STORAGE_PROPERTY_QUERY</a> structure.
     */
    public static class STORAGE_PROPERTY_QUERY extends Structure {
        STORAGE_PROPERTY_ID PropertyId = new STORAGE_PROPERTY_ID();
        STORAGE_QUERY_TYPE QueryType = new STORAGE_QUERY_TYPE();
        PrimitiveArray AdditionalParameters;

        public STORAGE_PROPERTY_QUERY() {
            AdditionalParameters = new PrimitiveArray(UInt8.class, 1);
            init(new Parameter[]{PropertyId, QueryType, AdditionalParameters}, (short) 4);
        }
    }

    /**
     * Wrapper for <a href = "http://msdn.microsoft.com/en-us/library/ff800835%28v=VS.85%29.aspx">STORAGE_DEVICE_DESCRIPTOR</a> structure.
     */
    public static class STORAGE_DEVICE_DESCRIPTOR extends Structure {
        static final int STORAGE_DEVICE_DESCRIPTOR_LENGTH = 48;
        UInt32 Version = new UInt32();
        UInt32 Size = new UInt32();
        UInt8 DeviceType = new UInt8();
        UInt8 DeviceTypeModifier = new UInt8();
        Bool RemovableMedia = new Bool();
        Bool CommandQueueing = new Bool();
        UInt32 VendorIdOffset = new UInt32();
        UInt32 ProductIdOffset = new UInt32();
        UInt32 ProductRevisionOffset = new UInt32();
        UInt32 SerialNumberOffset = new UInt32();
        STORAGE_BUS_TYPE BusType = new STORAGE_BUS_TYPE();
        UInt32 RawPropertiesLength = new UInt32();
        PrimitiveArray RawDeviceProperties;

        public STORAGE_DEVICE_DESCRIPTOR() {
            this(0);
        }

        public STORAGE_DEVICE_DESCRIPTOR(int bufferSize) {
            int propertiesLength;
            if (bufferSize == 0) {
                propertiesLength = 1;
            } else {
                propertiesLength = bufferSize - STORAGE_DEVICE_DESCRIPTOR_LENGTH;
            }

            RawDeviceProperties = new PrimitiveArray(Char.class, propertiesLength);
            init(new Parameter[]{
                    Version,
                    Size,
                    DeviceType,
                    DeviceTypeModifier,
                    RemovableMedia,
                    CommandQueueing,
                    VendorIdOffset,
                    ProductIdOffset,
                    ProductRevisionOffset,
                    SerialNumberOffset,
                    BusType,
                    RawPropertiesLength,
                    RawDeviceProperties
            }, (short) 4);
        }
    }

    public static void main(String[] args) throws Exception {
        char driveLetter = 'g';

        String drive = MessageFormat.format("\\\\.\\{0}:", new Object[] {new Char(driveLetter)});

        Handle volumeHandle = new Handle();
        Function createFile = Kernel32.getInstance().getFunction(new FunctionName("CreateFile").toString());
        long errorCode = createFile.invoke(volumeHandle, new Parameter[]{
                new Str(drive),
                new UInt32(0),
                new UInt32(FILE_SHARE_READ | FILE_SHARE_WRITE),
                new Pointer.Void(),
                new UInt32(OPEN_EXISTING),
                new UInt32(),
                new Pointer.Void()
        });

        if (volumeHandle.getValue() == Handle.INVALID_HANDLE_VALUE) {
            throw new LastErrorException(errorCode, "Failed to open drive: " + drive);
        }

        Function deviceIoControl = Kernel32.getInstance().getFunction("DeviceIoControl");
        Bool result = new Bool();

        STORAGE_PROPERTY_QUERY query = new STORAGE_PROPERTY_QUERY();
        query.PropertyId.setValue(STORAGE_PROPERTY_ID.StorageDeviceProperty);
        query.QueryType.setValue(STORAGE_QUERY_TYPE.PropertyStandardQuery);

        STORAGE_DEVICE_DESCRIPTOR deviceDescriptor = new STORAGE_DEVICE_DESCRIPTOR(1024);
        Pointer deviceDescriptorPtr = new Pointer(deviceDescriptor);

        UInt32 bytesReturned = new UInt32();

        try {
            errorCode = deviceIoControl.invoke(result, new Parameter[]{volumeHandle,
                    new UInt32(IOCTL_STORAGE_QUERY_PROPERTY),
                    new Pointer(query),
                    new UInt32(query.getLength()),
                    deviceDescriptorPtr,
                    new UInt32(deviceDescriptor.getLength()),
                    new Pointer(bytesReturned),
                    new Pointer.Void()
            });

            System.out.println("deviceDescriptor.BusType = " + deviceDescriptor.BusType.getValue());

            if (deviceDescriptor.BusType.getValue() == STORAGE_BUS_TYPE.BusTypeFileBackedVirtual) {
                System.out.println("Specified drive is virtual");
            } else {
                System.out.println("Specified drive is not virtual");
            }

            if (!result.getValue()) {
                throw new LastErrorException(errorCode, "Failed to get get IOCTL_STORAGE_QUERY_PROPERTY.");
            }
        } finally {
            if (!volumeHandle.isNull()) {
                Handle.closeHandle(volumeHandle);
            }
        }
    }
}
