#!/usr/bin/env python
# Copyright (c) 2020 Klustron inc. All rights reserved.
# This source code is licensed under Apache 2.0 License,
# combined with Common Clause Condition 1.0, as detailed in the NOTICE file.

import os
import os.path
import json
import argparse
import platform

def get_cpu_info():
    arch = platform.machine()
    if arch == 'aarch64':
        return ('aarch', '64')
    f = open('/proc/cpuinfo')
    str = f.read()
    f.close()
    has_vendor_id = False
    has_model_name = False
    vendor_id = None
    model_name = None
    for line in str.split('\n'):
        pair = line.split(':')
        if len(pair) < 2:
            continue
        name = pair[0].strip()
        value = pair[1].strip()
        if not has_vendor_id and name == 'vendor_id':
            vendor_id = value
            has_vendor_id = True
        if not has_model_name and name == 'model name':
            model_name = value
            has_model_name = True
        if has_vendor_id and has_model_name:
            break
    return (vendor_id, model_name)

def get_bios_info():
    f = open('/sys/class/dmi/id/uevent')
    str = f.read()
    f.close()
    uevent = None
    pair = str.split('=', 1)
    if len(pair) >= 2:
        uevent = pair[1].strip()
    return uevent

def get_disk_info():
    rets = []
    items = os.listdir('/sys/block')
    for item in items:
        path = '/sys/block/' + item + '/device/serial'
        if os.path.exists(path):
            f = open(path)
            str = f.read().strip()
            f.close()
            rets.append({'name': item, 'serial':str})
    return rets

if  __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Specify the arguments.')
    parser.add_argument('--output', type=str, help="The output path", default="machine_info.json")
    args = parser.parse_args()

    cpuinfo = get_cpu_info()
    biosinfo = get_bios_info()
    diskinfo = get_disk_info()
    f = open(args.output, 'w')
    mach_obj =  {
        "cpu_id": "-".join(cpuinfo),
        "bios_info": biosinfo,
        "device_infos": str(diskinfo)
            }
    json.dump(mach_obj, f, indent=4)
    f.close()
