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
import ConfigParser
from MythTV import MythDB, MythBE, Recorded, MythError
from socket import timeout, gethostname


storage_dir = "/etc/storage.d"
pickle_file = "%s/storage.pkl" %storage_dir


SG_MAP_BE={ 'Default'    :'tv',
'LiveTV'     :'media/tv/live',
'DB Backups' :'system_backups/mythtv_db/'}


SG_MAP_FE={
  'Videos'     :'media/video',
  'Trailers'   :'media/video_stuff/trailers',
  'Coverart'   :'media/video_stuff/coverart',
  'Fanart'     :'media/video_stuff/fanart',
  'Banners'    :'media/video_stuff/banners',
  'Screenshots':'media/video_stuff/screenshots',
  }

FS_LIST_BE=[]
for key in SG_MAP_BE.keys():
    FS_LIST_BE.append(SG_MAP_BE[key])

FS_LIST_FE=[]
for key in SG_MAP_FE.keys():
    FS_LIST_FE.append(SG_MAP_FE[key])


class disk_device:
    def  __init__(self,device,storage_dir):
        device_obj = bus.get_object("org.freedesktop.UDisks", device)
        device_props = dbus.Interface(device_obj, dbus.PROPERTIES_IFACE)
        self.storage_dir = storage_dir
        self.top_mount_dir = "/storage"
        self.config = ConfigParser.RawConfigParser()
        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.mmount = False
        self.f = 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()
        self.uuid=''
        self.new_mount_point=''

    def set_partition(self,partition):
        self.block_partition = "%s3" %self.block_path

    def set_mmount(self,flag):
        self.mmount = flag

    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_name(self):
        filename="%s_%s" %(self.model.replace(' ',''),
                                self.serial_number.replace(' ',''))
        return filename

    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)
        self.uuid = uuid.strip()
        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,bind=False):
        new_fstab_list=['UUID=', 'mount_point', 'auto', 'defaults', '0', '1']
        fstab=self.read_fstab()


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

        if bind:
            new_fstab_list=["/myth" , self.new_mount_point , "none" , "rw,bind", '0', '0']
        else:
        #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,mount=True):
        try:
            os.stat(self.new_mount_point)
        except:
            os.makedirs(self.new_mount_point)
        if mount == True:
            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,weight='0',install_call=False):
        print "    Adding to storage groups"
        sgweight=int(weight)
        for key in SG_MAP.iterkeys():
            #print key," : ", SG_MAP[key]
            gn=key
            hn=host
            dn="%s/%s" %(self.new_mount_point,SG_MAP[key])
            #print dn
            #print gn
            #print hn
            if install_call == True :
                print "Will write SG for stuff after the fact"
            else:
                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

                    if sgweight > 0:
                        try:
                            #("SGweightPerDir:server2:/mnt/video", 99, "server2");
                            sgw="SGweightPerDir:%s:%s" %(hn,dn)
                            #print sgw
                            #print sgweight
                            #print hn

                            c.execute("""insert into settings (value,data,hostname) values (%s,%s,%s)""",(sgw,sgweight,hn))
                            print "        Adding storage group weight of %s for %s\n" %(sgweight,gn)
                        except:
                            print "        *Error setting storage group weight %s for %s\n" %(sgweight,gn)

        return

    def write_config(self):
        print "    Writing out storage.d conf file"
        self.config.add_section('storage')
        self.config.set('storage','uuid',self.uuid)
        self.config.set('storage','mountpoint',self.new_mount_point)
        self.config.set('storage','shareable','True')
        self.config.set('storage','mmount',self.mmount)


        filename="%s_%s.conf" %(self.model.replace(' ',''),
                                self.serial_number.replace(' ',''))

        configfile="/etc/storage.d/%s" %filename
        print "     %s" %configfile
        with open(configfile, 'wb') as configfile:
            self.config.write(configfile)
        return

    def symlink(self):
        print "    Creating symlink"
        cmd = "ln -s %s %s/myth " %(self.new_mount_point, self.top_mount_dir)
        runcmd(cmd)
        cmd = "ln -s %s %s/data " %(self.new_mount_point, self.top_mount_dir)
        runcmd(cmd)



#end of class


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,storage_dir)

        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 "\n*   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("\n    Adding the drive will remove all contents on the drive. \n    Do 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 prompt_to_continue():
    loop = True
    #while loop:
    str1 = raw_input("\n    Ready to start, press Y to continue, anything else to abort:")

        #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, destruction, no_mount, install_call , dir_fe_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 "\n"
            print "    Storage is already in use or previously skipped:%s" %i.model
            print "        location: %s" %i.block_path
            print "        size: %s" %i.device_size
            continue

        else:
            if not scan_only:
                print "\n"
                print "-------------------------------------------------------------"
                print "    Found new hard drive: %s" %i.model
                print "    location: %s" %i.block_path
                print "    size: %s " %i.device_size

                if prompt_to_add(i) :
                    print "\n    Disk will be added to the storage pool!"
                    process_list.append(i)
            else:
                process_list.append(i)
    print "    End of scan"
    print "---------------------------------------"
    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)
    if len(process_list) > 0:
        print "\n    Adding storage: \n"
        if prompt_to_continue() == True:
            for i in process_list:
                print "    Drive: %s" %(i.get_name())
                if destruction == True:
                    i.partition_disk()
                    i.format_disk()
                i.add_fstab()
                i.mount_disk(no_mount)
                if destruction == True:
                    if dir_fe_only != True:
                        i.mkdirs(FS_LIST_BE)
                    i.mkdirs(FS_LIST_FE)
                i.write_config()

                #always create the FE groups(video)..do not need to create Be
                if dir_fe_only != True:
                    i.add_sg(DB,host,SG_MAP_BE)

                i.add_sg(DB,host,SG_MAP_FE)



                print "-----"

        #i.add_sg(DB,host,SG_MAP)

def myth_main(no_mount,install_call,dir_fe_only):
    global bus
    bus = dbus.SystemBus()
    #search for root
    f = open('/etc/fstab', 'r')
    fstab=f.readlines()
    f.close()
    for i in fstab:
        split_line=i.split()
        if not split_line:
            continue

        try:
            if split_line[1] ==  "/" :
                uuid_device = split_line[0]
                break
        except:
            print "Couldn't find / in fstab"
            sys.exit(1)
    if uuid_device.startswith("UUID"):
        uuid_device = uuid_device.split("=")[1]
        cmd = "blkid -U %s" %uuid_device
        device = runcmd(cmd)[1]
    else:
        device = uuid_device.strip()
    #should have something like /dev/sda1
    #remove all the digits
    device = ''.join([letter for letter in device if not letter.isdigit()])


    system_drive_list = scan_system()
    for i in system_drive_list:
        if i.block_path == device:
            break
    else:
        print "Could'nt find root device in block list"
        sys.exit(1)
    if not install_call == True:
        DB = MythDB()
        host=gethostname()
    else:
        DB = None
        host=gethostname()


    print "    Drive: %s" %(i.get_name())
    i.set_mmount(True)
    i.set_partition("3")
    i.add_fstab(True)
    #if not install_call:
    i.mount_disk(no_mount)
    i.write_config()
    #always create the FE groups(video)..do not need to create Be
    if dir_fe_only != True:
        i.add_sg(DB,host,SG_MAP_BE,'99',install_call)

    i.add_sg(DB,host,SG_MAP_FE,'99',install_call)

    i.symlink()


def usage():
    help='''\n
    Add storage is designed to find and setup new disks for mythtv usage.
    It's a powerfull that could destroy data if not used correctly, so please be careful.
    Normal operations include (in this order):
        Partition the disk
        Format the disk
        Add disk to /etc/fstab
        Mount the disk
        Create the directory
        Write out the config file to /etc/storage.d
        Add new locations to mythtv storage groups

    options:

    --no_mount :  Do not mount the disk, only add it to /etc/fstab and create the dir.
    --no_destruction:  Will not  partition or format the disk
    --new_init :  Erase the list of new disks and rescan.
    --report :  will scan the disks and print out if it found new storage.
    --fe_only :  Will only create the storage group dir for videos..excludes tv
    '''
    print help
    sys.exit(0)




if __name__ == "__main__":
    scan_only = False
    myth_mount = False
    no_mount = False
    destruction = True
    install_call = False
    dir_fe_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 "--help" in sys.argv:
        usage()

    if "--install_call" in sys.argv:
        install_call = True

    if "--no_mount" in sys.argv :
        no_mount = True

    if "--no_destruction" in sys.argv:
        destruction = False

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

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

    if "--fe_only" in sys.argv:
        dir_fe_only = True



    if "--double_myth" in sys.argv:
        myth_main(no_mount,install_call,dir_fe_only)
    else:
        main(scan_only,destruction,no_mount, install_call, dir_fe_only)