package main

// #include <stdlib.h>
import "C"

import (
    "bytes"
    "errors"
    "os"
    "syscall"
    "unsafe"
)

var (
    success                = "The operation completed successfully."
    kernel                 = syscall.NewLazyDLL("Kernel32.dll")
    setVolumeMountPoint    = kernel.NewProc("SetVolumeMountPointA")
    getValumeNameForMount  = kernel.NewProc("GetVolumeNameForVolumeMountPointA")
    deleteVolumeMountPoint = kernel.NewProc("DeleteVolumeMountPointA")
)

func main() {
    defer syscall.CloseHandle(syscall.Handle(kernel.Handle()))
    if len(os.Args) != 3 {
        usage()
    }
    if os.Args[1] == "-d" {
        let := os.Args[2]
        if len(let) != 1 {
            println("盘符不正确,确认盘符格式:-d D格式")
            return
        }
        let += ":\\"
        err := delVolumeMount(let)
        if err != nil {
            println(err.Error())
        }
        return
    }
    err := ChangeLetter(os.Args[1],os.Args[2])
    if err != nil {
        println(err.Error())
    }
}

//ChangeLetter("E","G")
func ChangeLetter(oldLetter,newLetter string) error {
    if len(oldLetter) > 1 || len(newLetter) != 1 {
        return errors.New("盘符不正确,确认盘符格式:C,D格式")
    }
    oldLetter += ":\\"
    newLetter += ":\\"
    volumeID,err := getValumeName(oldLetter)
    if err != nil {
        return err
    }
    println("获取到 volume id:",volumeID)
    err = delVolumeMount(oldLetter)
    if err != nil {
        return err
    }
    return setVolumeMount(newLetter,volumeID)
}

//getValumeName("D:\\")
func getValumeName(letter string) (string,error) {
    char := C.CString(letter)
    defer C.free(unsafe.Pointer(char))
    n := make([]byte, 254)
    _,_,err := getValumeNameForMount.Call(uintptr(unsafe.Pointer(char)),uintptr(unsafe.Pointer(&n[0])), 254)
    if err.Error() != "More data is available." {
        return "",err
    }
    return string(n[:bytes.Index(n,[]byte{0})]),nil
}

//setVolumeMount("G:\\",`\\?\Volume{3a7aa2c9-5dfc-11e6-99d9-806e6f6e6963}\`)
func setVolumeMount(letter,volumeName string) error {
    path := C.CString(letter)
    defer C.free(unsafe.Pointer(path))
    v := C.CString(volumeName)
    defer C.free(unsafe.Pointer(v))
    _,err := setVolumeMountPoint.Call(uintptr(unsafe.Pointer(path)),uintptr(unsafe.Pointer(v)))
    if err.Error() != success {
        return err
    }
    return nil
}

//delVolumeMount("D:\\")
func delVolumeMount(letter string) error {
    p := C.CString(letter)
    defer C.free(unsafe.Pointer(p))
    _,err := deleteVolumeMountPoint.Call(uintptr(unsafe.Pointer(p)))
    if err.Error() != success {
        return err
    }
    return nil
}

func usage() {
    println(os.Args[0] + ` usage: 删除一个盘符:` + os.Args[0] + " -d E" + ` 修改一个盘符:` + os.Args[0] + " D E")
    os.Exit(1)
}