#!/usr/bin/python3

# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

# Provides a list of physical block devices installed in
# the system. A block device can be a SATA SSD or NVMe drive.
import struct
import sys
from nv_subps import SubPs

DISC_HDR_LEN        = 48
TPER_FEATURE_LEN    = 16
LOCKING_FEATURE_LEN = 16

TPER_FEATURE_CODE    = 1
LOCKING_FEATURE_CODE = 2

SECURITY_RECV       = 0x82

class Level0Discovery:

    def __init__(self):

        self.hdr_param_len   = 0    # Parameter data length
        self.hdr_struct_rev  = 0    # Data structure rev
        self.hdr_vs          = 0    # Vendor specific

        self.tper_feature_code = -1
        self.tper_version     = 0
        self.tpter_len        = 0

        self.locking_feature_code = -1
        self.locking_version      = 0
        self.locking_supported    = 0
        self.locking_enabled      = 0
        self.locked               = 0

    def is_sed_aware(self):

        if self.tper_feature_code == TPER_FEATURE_CODE and \
            self.locking_feature_code == LOCKING_FEATURE_CODE:
            return True

        return False

    def print_level0_disc_header(self):
        print("Level0 Discovery Header %x: " %(self.hdr_param_len))
        print("Data Structure Rev %x: " %(self.hdr_struct_rev))
        print("Vendor Specific %x: " %(self.hdr_vs))

    def print_level0_disc_tper_feature(self):
        print("TPER Feature Code %x: " %(self.tper_feature_code))
        print("TPER Version: %x" %(self.tper_version))
        print("TPER Length %x: " %(self.tper_len))

    def print_level0_locking_feature(self):
        print("Locking Feature Code %x: " %(self.locking_feature_code))
        print("Locking Version: %x" %(self.locking_version))
        print("Locking Supported %x: " %(self.locking_supported))
        print("Locking Enabled: %x" %(self.locking_enabled))
        print("Locked %x: " %(self.locked))

    def process_disc_resp_data(self,resp_buf):

        self.hdr_param_len, self.hdr_struct_rev, _, \
        self.hdr_vs = struct.unpack_from('>IIII', resp_buf, 0)

        self.tper_feature_code, self.tper_version, \
        self.tper_len, _ = struct.unpack_from('>HBBB', resp_buf, DISC_HDR_LEN)
        self.tper_version = self.tper_version >> 4

        self.locking_feature_code, self.locking_version, \
        _, tmp = struct.unpack_from('>HBBB', resp_buf, DISC_HDR_LEN + TPER_FEATURE_LEN)

        self.locking_version = self.locking_version >> 4
        self.locking_supported = (tmp & 1)
        self.locking_enabled = ((tmp >> 1) & 1)
        self.locked = ((tmp >> 2) & 1)

def test():
    l0_disc = Level0Discovery()
    protocol = 1
    comID = 1
    cdw10 = protocol << 24 | comID << 8
    buflen = 800

    sps = SubPs()
    cmd = "nvme admin-passthru -o 0x%x -4 0x%x -5 %d -l %d %s -r -b" \
            %(SECURITY_RECV, cdw10, buflen, buflen, sys.argv[1])

    rc = sps.run_cmd(cmd.split(' '))
    if rc != 0:
        print("cmd failed with return code %x" %(rc))
        exit(1)

    data = sps.output.encode('utf-8')
    l0_disc.process_disc_resp_data(data)
    l0_disc.print_level0_disc_header()
    l0_disc.print_level0_disc_tper_feature()
    l0_disc.print_level0_locking_feature()