path: root/abs/core/LinHES-system/add_storage.py

#!/usr/bin/python2
#Program used to auto_add new storage to mythtv storage groups
#If it's a new disk it will erase the entire disk and reformat
#
#Drives that are mount, in fstab, size < 5000 bytes, optical or have already been seen will not be presented as an option.
#


import dbus
import pickle
import commands
import sys,os
import random, string
from MythTV import MythDB, MythBE, Recorded, MythError
from socket import timeout, gethostname



pickle_file="/usr/LH/etc/storage.pkl"



SG_MAP={ 'Default'    :'tv',
  'LiveTV'     :'live_tv',
  'Videos'     :'video',
  'Trailers'   :'video/trailers',
  'Coverart'   :'video/converart',
  'Fanart'     :'video/fanart',
  'Banners'    :'video/banners',
  'Screenshots':'video/screenshots',
  'DB Backups' :'db_backups'}
FS_LIST=[]
for key in SG_MAP.keys():
    FS_LIST.append(SG_MAP[key])



class disk_device:
    def  __init__(self,device):
        device_obj = bus.get_object("org.freedesktop.UDisks", device)
        device_props = dbus.Interface(device_obj, dbus.PROPERTIES_IFACE)
        self.fs_map = self.get_fsmap()
        self.is_device = self.get_is_device(device_props)
        self.vendor = self.get_vendor(device_props)
        self.model = self.get_model(device_props)

        self.mount_path = self.get_mounts(device_props)
        self.is_mounted = self.get_is_mounted(device_props)

        self.parition_size = self.get_partition_size(device_props)
        self.device_size = self.get_device_size(device_props)
        self.serial_number = self.get_serial_number(device_props)
        self.read_only = self.get_is_readonly(device_props)
        self.is_optical = self.get_is_optical_disc(device_props)
        self.connection = self.get_connections(device_props)
        self.block_path = self.get_device_file(device_props)
        self.device_file_path = self.get_device_file_path(device_props)
        self.block_partition="%s1" %self.block_path
        self.in_use = self.get_in_use()


    def get_is_device(self,device_props):
        return device_props.Get('org.freedesktop.UDisks.Device', "DeviceIsDrive")

    def get_vendor(self,device_props):
        return device_props.Get('org.freedesktop.UDisks.Device', "DriveVendor")

    def get_model(self,device_props):
        return device_props.Get('org.freedesktop.UDisks.Device', "DriveModel")

    def get_mounts(self,device_props):
        return device_props.Get('org.freedesktop.UDisks.Device', "DeviceMountPaths")

    def get_is_mounted(self,device_props):
       return device_props.Get('org.freedesktop.UDisks.Device', "DeviceIsMounted")

    def get_in_use(self):
        in_use = False
        for i in self.fs_map:
            if self.block_path in i[0]:
                in_use = True
                break
        return in_use


    def get_partition_size(self,device_props):
        return device_props.Get('org.freedesktop.UDisks.Device', "PartitionSize")

    def get_device_size(self,device_props):
        return device_props.Get('org.freedesktop.UDisks.Device', "DeviceSize")


    def get_serial_number(self,device_props):
        serial_number = device_props.Get('org.freedesktop.UDisks.Device', "DriveSerial")
        random_string = os.urandom(5)
        if serial_number == '':
            serial_number = "".join( [random.choice(string.letters) for i in xrange(6)] )
            serial_number = "NSW%s" %serial_number
        return serial_number

    def get_is_readonly(self,device_props):
        return device_props.Get('org.freedesktop.UDisks.Device', "DeviceIsReadOnly")

    def get_is_optical_disc(self,device_props):
        return device_props.Get('org.freedesktop.UDisks.Device', "DeviceIsOpticalDisc")

    def get_connections(self,device_props):
       return device_props.Get('org.freedesktop.UDisks.Device', "DriveConnectionInterface")

    def get_device_file(self,device_props):
       return  device_props.Get('org.freedesktop.UDisks.Device', "DeviceFile")

    def get_device_file_path(self,device_props):
        path = device_props.Get('org.freedesktop.UDisks.Device', "DeviceFileByPath")
        try:
            path = path[0]
        except:
            path = "None"

        return path


    def partition_disk(self):
        print "    Creating new partiton table"
        cmd = "parted -s -a optimal %s mklabel gpt" %self.block_path
        runcmd(cmd)
        cmd = "parted -s -a optimal %s mkpart primary \" 1 -1\"" %self.block_path
        runcmd(cmd)
        return

    def get_fsmap(self):
        #block,point,fs
        fs_map=[]

        f = open('/proc/mounts','r')
        mounts=f.readlines()
        f.close()
        for line in mounts:
            temp_fs=[]
            split_line=line.split()
            block=split_line[0]
            mountp=split_line[1]
            fs=split_line[2]
            block=os.path.realpath(block)
            if block.startswith("/dev"):
                temp_fs.append(block)
                temp_fs.append(mountp)
                temp_fs.append(fs)
                fs_map.append(temp_fs)
        return fs_map

    def find_fstype(self,moutpoint):
        fstype="ext3"
        mp=['/myth', '/data/media']
        for i in self.fs_map:
                if i[1] in mp:
                    fstype = i[2]
                    break
        return fstype


    def format_disk(self):
        fstab = self.read_fstab()
        #lookup format
        current_media_mount = self.find_options_type(fstab)[1]
        new_fstype = self.find_fstype(current_media_mount)

        #do format
        cmd = "mkfs -t %s %s " %(new_fstype,self.block_partition)
        print "    Formating %s with %s" %(self.block_partition,new_fstype)
        runcmd(cmd)
        return

    def find_uuid(self,partition):
        #logging.info("Finding the UUID for %s...", partition)
        cmd = "blkid -s UUID %s" %partition
        tmpuuid = runcmd(cmd)[1]
        splituuid = tmpuuid.partition("=")
        uuid = splituuid[2].replace('"', "")
       # logging.info("The uuid is %s", uuid)
        if uuid == '':
            print "Could not find a UUID for device: %s" %partition
            sys.exit(1)

        return uuid.strip()

    def read_fstab(self):
        f = open('/etc/fstab', 'r')
        fstab=f.readlines()
        f.close()
        return fstab

    def find_options_type(self,fstab):
        mp=['/myth', '/data/media']
        for i in fstab:
            split_line=i.split()
            try:
                if split_line[1] in mp:
                    options = split_line[3]
                    break
                else:
                    options = "defaults"
                    mount_point = i
            except:
                options = "defaults"
        return options,i


    def add_fstab(self):
        new_fstab_list=['UUID=', 'mount_point', 'auto', 'defaults', '0', '1']
        fstab=self.read_fstab()

        #determine mount_path
        self.new_mount_point="/data/media/%s_%s" %(self.model.replace(' ',''),self.serial_number.replace(' ',''))

        #determine options
        new_options = self.find_options_type(fstab)[0]

        #find blkid
        self.block_partition="%s1" %self.block_path
        uuid=self.find_uuid(self.block_partition)

        #construct new line
        new_fstab_list[0]="UUID=%s" %uuid
        new_fstab_list[1]=self.new_mount_point
        new_fstab_list[3]=new_options

        new_fstab_line='\t'.join(new_fstab_list)
        new_fstab_line="%s\n" %new_fstab_line
        fstab.append(new_fstab_line)

        #add line to fstab
        f = open('/etc/fstab', 'w')
        #f = open('/tmp/fstab', 'w')
        for i in fstab:
            f.write(i)
            #f.write("\n")
        f.close()
        return

    def mount_disk(self):
        try:
            os.stat(self.new_mount_point)
        except:
            os.makedirs(self.new_mount_point)
        cmd = "mount %s" %self.new_mount_point
        runcmd(cmd)
        return

    def mkdirs(self,FS_LIST):
        print "    Creating Directory stucture"
        for y in FS_LIST:
            new_dir="%s/%s" %(self.new_mount_point,y)
            try:
                os.stat(new_dir)
            except:
                os.makedirs(new_dir)
                cmd="chown -R mythtv:mythtv /%s" %self.new_mount_point
                runcmd(cmd)
                cmd="chmod -R 775 /%s" %self.new_mount_point
                runcmd(cmd)

    def add_sg(self,DB,host,SG_MAP):
        print "    Adding to storage groups"
        for key in SG_MAP.iterkeys():
            #print key," : ", SG_MAP[key]
            #try:
                gn=key
                hn=host
                dn="%s/%s" %(self.new_mount_point,SG_MAP[key])
                #print dn
                #print gn
                #print hn
                with DB as c:
                    try:
                        c.execute("""insert into storagegroup (groupname,hostname,dirname) values (%s,%s,%s)""",(gn,hn,dn))
                        print "        Adding  location: %s  to storagegroup %s" %(dn,gn)
                    except:
                        print "        Error inserting %s into storage groups" %dn
        return



def runcmd(cmd):
    if True :
        pass
    else:
        cmd = "echo "+cmd
    #print cmd
    cmdout = commands.getstatusoutput(cmd)
    #logging.debug("    %s", cmdout)
    return cmdout




def scan_system():

    ud_manager_obj = bus.get_object("org.freedesktop.UDisks", "/org/freedesktop/UDisks")
    ud_manager = dbus.Interface(ud_manager_obj, 'org.freedesktop.UDisks')
    current_drive_list=[]
    for dev in ud_manager.EnumerateDevices():
        drive = disk_device(dev)

        if drive.is_device and drive.device_size > 5000 and not drive.is_optical :
            current_drive_list.append(drive)
    return current_drive_list

def read_known_list():
    #reading pickle file
    known_drive_list=[]
    try:
        pkl_file = open(pickle_file, 'rb')
        known_drive_list = pickle.load(pkl_file)
        pkl_file.close()
    except:
        pass
    return known_drive_list


def write_known_drive_list(known_drive_list):
    output = open(pickle_file, 'wb')
    pickle.dump(known_drive_list, output, -1)
    output.close()


def search_for_match(system_drive,known_drive_list):
    match_drive=False
    for y in known_drive_list:
        if system_drive.serial_number.startswith("NSW"):
            #print "Match_test: hash"
            system_drive_hash = "%s_%s_%s" %(system_drive.model,system_drive.parition_size,system_drive.device_file_path)
            y_drive_hash = "%s_%s_%s" %(y.model,y.parition_size,y.device_file_path)
            if system_drive_hash == y_drive_hash :
                match_drive = True
                print "No serial number was found, matched using hash method: %s" %system_drive.model
                break

        elif y.serial_number == system_drive.serial_number:
            #print "Match_test: serial number"
            match_drive=True
            break

    return match_drive


def prompt_to_add(current_drive):
    loop = True
    while loop:
        str1 = raw_input("     Adding the drive will remove all contents on the drive. \n\nDo you wish enable this drive for MythTV storage(Y/N)?:")

        if str1 in ['Y','N','y','n']:
            loop = False
            break
        print "\n"
    if str1 == 'Y' or str1 == 'y':
        rc = True
    else:
        rc = False
    return rc

def remove_pickle():
    try:
        print "Resetting list of known drives."

        os.remove(pickle_file)
    except:
        pass


#--------------------------------------------

def main(scan_only):
    global bus
    bus = dbus.SystemBus()

    system_drive_list = scan_system()
    known_drive_list=[]
    known_drive_list = read_known_list()
    process_list=[]



    for i in system_drive_list:
        #print i.mount_path
        #print i.is_mounted
        #print i.in_use
        #print i.model
        #print i.block_path
        #print "--"

        if search_for_match(i,known_drive_list) or i.in_use :
            print "    Storage is already in use or previously skipped: %s location: %s size: %s" %(i.model,i.block_path,i.device_size)
            continue

        else:
            print "\n"
            if not scan_only:
                print "\n"
                print "-------------------------------------------------------------"

                print "Found new hard drive: %s location: %s size: %s \n" %(i.model,i.block_path,i.device_size)

                if prompt_to_add(i) :
                    print "\nDisk will be added to the storage pool!"
                    process_list.append(i)
            else:
                process_list.append(i)

    if scan_only:
        if len(process_list) > 0:
            f = open('/tmp/scan_report', 'w')
            for i in process_list:
                f.write("drive: %s , location: %s ,size: %s \n" %(i.model,i.block_path,i.device_size))
                print "drive: %s , location: %s ,size: %s \n" %(i.model,i.block_path,i.device_size)
            f.close()

        sys.exit(0)

    if len(process_list) > 0:
        DB = MythDB()
        host=gethostname()
    else:
        print "\nDid not find any new storage to add.\n"
            #BE = MythBE(db=DB)


    #save new list to disk_device
    write_known_drive_list(system_drive_list)

    for i in process_list:
        i.partition_disk()
        i.format_disk()
        i.add_fstab()
        i.mount_disk()
        i.mkdirs(FS_LIST)
        i.add_sg(DB,host,SG_MAP)




if __name__ == "__main__":
    scan_only = False
    try:
        os.remove("/tmp/scan_report")
    except:
        pass

    if not os.geteuid()==0:
        sys.exit("\nRoot access is required to run this program\n")

    if "--new_init" in sys.argv :
        remove_pickle()

    if "--report" in sys.argv :
        scan_only = True

    main(scan_only)