#!/usr/bin/env python
import os
import glob
import errno
import stat
import sys
import subprocess


def main(argv):
    if len(argv) > 0:
        nvme_devices = argv
    else:
        nvme_devices = glob.glob('/dev/nvme*n1')

    for nvme_device in nvme_devices:
        device_name = get_device_name(nvme_device)

        if device_name is not None:
            print('Symlink {} -> {}'.format(device_name, nvme_device))
            force_symlink(nvme_device, device_name)


def get_device_name(nvme_device):
    if not os.path.exists(nvme_device):
        print("Device {} doesn't exist".format(nvme_device))
        return None

    mode = os.stat(nvme_device).st_mode
    if not stat.S_ISBLK(mode):
        print('Device {} is not a block device'.format(nvme_device))
        return None

    vendor_data = subprocess.check_output(['nvme', 'id-ctrl', '-b', nvme_device])
    device_offset = vendor_data.find(b'\x2f\x64\x65\x76\x2f')
    end = vendor_data[device_offset:].find(b'\x00')
    device_name = vendor_data[device_offset:][:end].rstrip()

    if len(device_name) > 0:
        return device_name


def force_symlink(src, dst):
    try:
        os.symlink(src, dst)
    except OSError, e:
        if e.errno == errno.EEXIST:
            os.remove(dst)
            os.symlink(src, dst)
        else:
            raise e


if __name__ == '__main__':
   main(sys.argv[1:])