#!/usr/bin/python2
#add_storage.py used to auto add new storage to MythTV storage groups
#If it's a new disk it will erase the entire disk and reformat.
#
#Disks that are mounted, 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 ConfigParser import SafeConfigParser
import glob

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={
    'Default'    :'media/tv/',
    'LiveTV'     :'media/tv/live/',
    'DB Backups' :'backup/mythtv_backups/',
    'Music'      :'media/music/',
    'Streaming'  :'media/streaming/',
    'Videos'     :'media/video/',
    'Photographs':'media/photos/',
    'Banners'    :'media/artwork/banners/',
    'Coverart'   :'media/artwork/coverart/',
    'Fanart'     :'media/artwork/fanart/',
    'MusicArt'   :'media/artwork/musicart/',
    'Screenshots':'media/artwork/screenshots/',
    'Trailers'   :'media/artwork/trailers/',
    }

FS_LIST=[]
for key in SG_MAP.keys():
    FS_LIST.append(SG_MAP[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 = "/data/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.dir_sg = 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=''
        self.disk_num=''

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

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

    def set_dir_sg(self,flag):
        self.dir_sg = flag

    def set_disk_num(self,num):
        self.disk_num=num

    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 partition 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="xfs"
        mp=['/myth', '/data/storage/disk0']
        for i in self.fs_map:
                if i[1] in mp:
                    fstype = i[2]
                    break
        return fstype

    def lookup_format(self):
        fstab = self.read_fstab()
        current_media_mount = self.find_options_type(fstab)[1]
        new_fstype = self.find_fstype(current_media_mount)
        #setting self.new_fstype so that it can be referenced when setting fstab
        self.new_fstype = new_fstype
        return

    def format_disk(self):
        #lookup format
        #self.lookup_format()
        #do format
        if self.new_fstype == "xfs":
            cmd = "mkfs -t %s -f %s " %(self.new_fstype,self.block_partition)
        else:
            cmd = "mkfs -t %s %s " %(self.new_fstype,self.block_partition)
        print "    Formatting %s with %s" %(self.block_partition,self.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/storage/disk0']
        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()
        new_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=["/data/storage/disk0" , self.new_mount_point , "none" , "rw,bind", '0', '0']
            uuid=self.find_uuid(self.block_partition)
        else:
            #check for old mount point and comment out
            for line in fstab:
                if not line.startswith("#"):
                    if line.find(self.new_mount_point) > -1:
                        print "      Found old mount %s in fstab, commenting out" %self.new_mount_point
                        line = "#"+line
                new_fstab.append(line)
            fstab=new_fstab

            #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,no_mount=False):
        try:
            os.stat(self.new_mount_point)
        except:
            os.makedirs(self.new_mount_point)
        if no_mount == False:
            if os.path.ismount(self.new_mount_point):
                print "    Disk already mounted, will not mount:\n       %s" %self.new_mount_point
                pass
            else:
                print "    Mounting %s" %self.new_mount_point
                cmd = "mount %s" %self.new_mount_point
                runcmd(cmd)
        return

    def mkdirs(self,FS_LIST):
        print "    Creating directory structure:"
        print "       %s" %self.new_mount_point
        for y in FS_LIST:
            print "          %s" %y
            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 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:
                    #delete old dir without trailing slash
                    c.execute("""delete from storagegroup where groupname = %s and hostname = %s and dirname = %s""", (gn,hn,dn.rstrip('/')))

                    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

                            #delete old dir without trailing slash
                            c.execute("""delete from settings where value = %s and data = %s and hostname = %s""", (sgw.rstrip('/'),sgweight,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 /etc/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','fstype',self.new_fstype)
        self.config.set('storage','shareable','True')
        self.config.set('storage','mmount',self.mmount)
        self.config.set('storage','storage_groups',self.dir_sg)
        self.config.set('storage','disk_num',self.disk_num)


        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_disk(self):
        print "    Creating symlink for disk%s" %self.disk_num
        disk_ln="%s/disk%s" %(self.top_mount_dir,self.disk_num)
        cmd = "ln -s %s %s" %(self.new_mount_point,disk_ln)
        runcmd(cmd)

    def symlink(self):
        pass
        print "    Creating symlink for /myth"
        if not os.path.exists("/myth"):
            cmd = "ln -s %s/media /myth " %(self.new_mount_point)
            runcmd(cmd)
        else:
            print "       Skipping symlink, /myth already exists"


#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,destruction = True):
    loop = True
    if destruction :
        prompt = '''
        ** Adding this disk will remove all contents on the disk. **
        This disk will be partitioned and formatted.

        Enable this disk for storage (Y/N)?:'''
    else:
        prompt = '''
        ** Preserving existing contents on the disk. **
        This disk will NOT be partitioned or formatted.

        Enable this disk for storage (Y/N)?:'''
    while loop:
        str1 = raw_input(prompt)

        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(process_list):
    loop = True
    print "\n\n\n   Ready to add additional storage!\n"
    if destruction:
        print "** WARNING: These disk(s) WILL be partitioned and formatted. **\n   ** All content on these disk(s) will be erased. **"
    else:
        print "   ** These disk(s) will NOT be partitioned and formatted. **"
    for i in process_list:
        print "      %s" %(i.get_name())
    str1 = raw_input("\n   Press Y to add disk(s), any other key to cancel:")

    if str1 == 'Y' or str1 == 'y':
        rc = True
    else:
        rc = False
        print "\nCancelled: No disk(s) added to your system."
    print "-----"
    return rc

def prompt_sg(dir_sg):
    #check for storage groups
    print "*" * 60
    if dir_sg != True:
        loop = True
        prompt_string='''
    MythTV Storage Groups are used for artwork, database backups,
    photos, music, streaming, TV recordings, and videos.

    The content on these storage groups will
    only be available while the system is online.

    Enabling MythTV Storage Groups will create the directories
    on the disk(s) and add the paths to the MythTV database.

    Enable MythTV Storage Groups (Y/N)?:'''

        while loop:
            str1 = raw_input(prompt_string)
            if str1 in ['Y','N','y','n']:
                loop = False
                break
            print "\n"

        if str1 == 'Y' or str1 == 'y':
            dir_sg = True
            print "    ** Will add MythTV Storage Groups!"
        else:
            print "    ** Will NOT add MythTV Storage Groups!"
            dir_sg = False
    else:
        dir_sg = True
        print "\n    --add_sg option used"
        print "    ** Will add MythTV Storage Groups!"

    return dir_sg

def remove_pickle():
    try:
        print "\n* Removing list of known disks.\n"
        os.remove(pickle_file)
    except:
        pass

def last_disk_num():
    parser = SafeConfigParser()
    num_list = []
    for conf_file in glob.glob('%s/*.conf' %storage_dir):
        parser.read(conf_file)
        disk_num = parser.get('storage', 'disk_num')
        num_list.append(int(disk_num))
        num_list.sort()
    return num_list[-1]
#--------------------------------------------

def main(scan_only, destruction, no_mount, install_call, dir_sg):
    global bus
    bus = dbus.SystemBus()

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

    print "-" * 60
    print "  Scan for Disks"

    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 :
            if search_for_match(i,known_drive_list) :
                dstatus = "    Ignoring - Disk has been previously skipped:"
            if i.in_use :
                dstatus = "    Ignoring - Disk is mounted:"
            if search_for_match(i,known_drive_list) and i.in_use :
                dstatus = "    Ignoring - Disk has been previously skipped and is mounted:"

            print "\n"
            print "    --------------------------------------------------------"
            print dstatus
            print "        model: %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 Disk:"
                print "        model: %s" %i.model
                print "        location: %s" %i.block_path
                print "        size: %s " %i.device_size

                if prompt_to_add(i,destruction) :
                    print "\n    %s will be added to your system!" %i.model
                    process_list.append(i)
                else:
                    no_process_list.append(i)
            else:
                process_list.append(i)
    print "\n"
    print "  Scan Finished"
    print "-" * 60

    if scan_only:
        if len(process_list) > 0:
            print "    Unknown or Unmounted Disks:"
            f = open('/tmp/scan_report', 'w')
            for i in process_list:
                f.write("disk: %s , location: %s ,size: %s \n" %(i.model,i.block_path,i.device_size))
                print "\n"
                print "    ---------------------------------------------------------"
                print "    Found New Disk:"
                print "        model: %s" %i.model
                print "        location: %s" %i.block_path
                print "        size: %s " %i.device_size
            f.close()
        sys.exit(0)

    for i in no_process_list:
        system_drive_list.append(i)


    if len(process_list) > 0:
        DB = MythDB()
        host=gethostname()
        for y in process_list:
            system_drive_list.remove(y)
        write_known_drive_list(system_drive_list)
    else:
        print "\nThere are no disks to add to your system.\n\nFor more options: add_storage.py --help\n"
        write_known_drive_list(system_drive_list)
            #BE = MythBE(db=DB)


    #save new list to disk_device
    #write_known_drive_list(system_drive_list)
    #disk_num = last_disk_num()


    if len(process_list) > 0:
        print "\n  Will add %s disk(s) to your system." %len(process_list)

        dir_sg = prompt_sg(dir_sg)
        if prompt_to_continue(process_list) == True:
            write_known_drive_list(system_drive_list)
            disk_num = last_disk_num()
            for i in process_list:
                print "    Disk: %s" %(i.get_name())
                disk_num = disk_num + 1
                i.lookup_format()
                if destruction == True:
                    i.partition_disk()
                    i.format_disk()
                i.add_fstab()
                i.mount_disk(no_mount)
                #if destruction == True:

                if dir_sg == True:
                    i.mkdirs(FS_LIST)

                i.set_disk_num(disk_num)
                i.set_dir_sg(dir_sg)
                i.write_config()
                system_drive_list.append(i)
                write_known_drive_list(system_drive_list)

                i.symlink_disk()

                if dir_sg == True:
                    i.add_sg(DB,host,SG_MAP)

                print "-----"
            cmd = "systemconfig.py -m fileshare"
            runcmd(cmd)
        #i.add_sg(DB,host,SG_MAP)

def myth_main(no_mount,install_call,dir_sg):
    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 "Couldn't find root device in block list"
        sys.exit(1)
    if not install_call == True:
        DB = MythDB()
        host=gethostname()
    else:
        DB = None
        host=gethostname()


    print "    Disk: %s" %(i.get_name())
    i.set_mmount(True)
    i.set_dir_sg(dir_sg)
    i.set_partition("7")
    i.set_disk_num(0)
    i.lookup_format()
    i.add_fstab(True)
    #if not install_call:
    i.mount_disk(no_mount)
    i.write_config()
    #no need to make the sub directories because the install process has taken care of it.
    if dir_sg == True:
        i.add_sg(DB,host,SG_MAP,'99',install_call)

    i.symlink()
    cmd = "systemconfig.py -m fileshare"
    runcmd(cmd)

def reconstruct_storagegroups():
    print "\nRecreating Storage Groups from contents of /etc/storage.d/\n"

    DB = MythDB()
    host=gethostname()

    for conf_file in glob.glob('%s/*.conf' %storage_dir):
        parser = SafeConfigParser()
        parser.read(conf_file)
        mount_point = parser.get('storage', 'mountpoint')
        mmount = parser.getboolean('storage', 'mmount')
        if not os.path.ismount(mount_point):
            print "Skipping: " + mount_point + " not mounted"
            continue
        try:
            dir_sg = parser.getboolean('storage', 'storage_groups')
        except ConfigParser.NoOptionError, err:
            print "SG not found in conf, get setting from DB"
            dir_sg = False
            # Get Storage Group directories from DB
            recs = DB.getStorageGroup(groupname="Default")
            for record in recs:
                if record.dirname.startswith(mount_point):
                    dir_sg = True
            # Write SG usage to conf
            parser.set('storage','storage_groups',str(dir_sg))
            with open(conf_file, 'wb') as conf_file:
                parser.write(conf_file)

        if dir_sg is True:
            print "SGs Enabled for: " + mount_point
            print "    Creating directory structure:"
            print "       %s" %mount_point
            for y in FS_LIST:
                new_dir="%s/%s" %(mount_point,y)
                try:
                    os.stat(new_dir)
                    print "          %s - exists" %y
                except:
                    os.makedirs(new_dir)
                    cmd="chown -R mythtv:mythtv /%s" %new_dir
                    runcmd(cmd)
                    cmd="chmod -R 775 /%s" %new_dir
                    runcmd(cmd)
                    print "          %s - created" %y

            print "    Adding storage groups to DB"
            if mmount is True:
                sgweight=99
            else:
                sgweight=0

            for key in SG_MAP.iterkeys():
                gn=key
                hn=host
                dn="%s/%s" %(mount_point,SG_MAP[key])
                with DB as c:
                    #delete old dir without trailing slash
                    c.execute("""delete from storagegroup where groupname = %s and hostname = %s and dirname = %s""", (gn,hn,dn.rstrip('/')))

                    try:
                        c.execute("""insert into storagegroup (groupname,hostname,dirname) values (%s,%s,%s)""",(gn,hn,dn))
                        print "        Added: %s to storagegroup %s" %(dn,gn)
                    except:
                        print "        Skipping: %s exists" %dn
                    if sgweight > 0:
                        try:
                            sgw="SGweightPerDir:%s:%s" %(hn,dn)

                            #delete old dir without trailing slash
                            c.execute("""delete from settings where value = %s and data = %s and hostname = %s""", (sgw.rstrip('/'),sgweight,hn))

                            if DB.settings[hn][sgw] == u'99':
                                print "        Skipping: storage group weight DB entry exists"
                            else:
                                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)

        else:
            print "SGs Disabled for: " + mount_point
    return

class reconstruct_path:
    def  __init__(self,conf_file):
        self.conf_file = conf_file
        parser = SafeConfigParser()
        parser.read(self.conf_file)
        self.config = ConfigParser.RawConfigParser()

        self.uuid = parser.get('storage', 'uuid')
        self.mount_point = parser.get('storage', 'mountpoint')
        self.shareable = parser.get('storage', 'shareable')
        self.myth_mount = parser.get('storage', 'mmount')
        self.bind = self.myth_mount
        self.disk_num = parser.get('storage', 'disk_num')
        self.top_mount_dir = os.path.dirname(self.mount_point)
        try:
            self.fstype = parser.get('storage', 'fstype')
        except:
            self.fstype = self.get_fstype()

    def get_fstype(self):
        cmd = "fsck -N UUID=%s" %self.uuid
        tmpfstype = runcmd(cmd)
        tmpfstype = tmpfstype[1].split('/sbin/fsck.')
        tmpfstype = tmpfstype[1].split(' ')
        self.fstype = tmpfstype[0]
        self.write_config()
        return self.fstype

    def get_conf(self):
        return self.conf_file

    def get_uuid(self):
        return self.uuid

    def get_mount_point(self):
        return self.mount_point

    def get_shareable(self):
        return self.shareable

    def get_is_myth_mount(self):
        return self.myth_mount

    def get_disk_num(self):
        return self.disk_num

    def create_mount_point(self):
        try:
            os.stat(self.mount_point)
        except:
            os.makedirs(self.mount_point)

    def find_options_type(self,fstab):
        mp=['/myth', '/data/storage/disk0']
        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 read_fstab(self):
        f = open('/etc/fstab', 'r')
        fstab=f.readlines()
        f.close()
        return fstab

    def check_in_fstab(self,fstab,check_path):
        for line in fstab:
            if line.find(check_path) > -1:
                return True
        return False

    def append_fstab(self,line):
        new_fstab_line='\t'.join(line)
        new_fstab_line="%s\n" %new_fstab_line

        f = open('/etc/fstab', 'a')
        #f = open('/tmp/fstab', 'a')
        f.write(new_fstab_line)
        #f.write("\n")
        f.close()

    def symlink(self):
        print "    Creating symlink for /myth"
        if not os.path.exists("/myth"):
            cmd = "ln -s %s/media /myth " %(self.mount_point)
            runcmd(cmd)
        else:
            print "       Skipping symlink, /myth already exists"

    def symlink_disk(self):
        print "    Creating symlink for disk%s" %self.disk_num
        disk_ln="%s/disk%s" %(self.top_mount_dir,self.disk_num)
        cmd = "ln -s %s %s" %(self.mount_point,disk_ln)
        runcmd(cmd)

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

        if self.bind == "True":
            self.symlink()

        if self.check_in_fstab(fstab,self.uuid) == True:
            print "    Found UUID of disk in fstab, will not add it"
        else:
            print "    Adding storage to fstab"
            if self.bind == "True" :
                print "    Bind mount"
                new_fstab_list=["/data/storage/disk0" , self.mount_point , "none" , "rw,bind", '0', '0']
                if self.check_in_fstab(fstab,self.mount_point) == False:
                    self.append_fstab(new_fstab_list)
                else:
                    print "    Found bind storage in fstab, will not add it"

            else:
                #check for old mount point and comment out
                f = open('/etc/fstab', 'w')
                for line in fstab:
                    if not line.startswith("#"):
                        if line.find(self.mount_point) > -1:
                            print "      Found old mount %s in fstab, commenting out" %self.mount_point
                            line = "#"+line
                    f.write(line)
                f.close()

                #construct new line
                new_options = self.find_options_type(fstab)[0]
                new_fstab_list[0]="UUID=%s" %self.uuid
                new_fstab_list[1]=self.mount_point
                new_fstab_list[3]=new_options
                self.append_fstab(new_fstab_list)

    def mount_disk(self,no_mount=False):
        try:
            os.stat(self.mount_point)
        except:
            os.makedirs(self.mount_point)
        if no_mount == False :
            if os.path.ismount(self.mount_point):
                print "    Disk already mounted, will not mount:\n       %s" %self.mount_point
                pass
            else:
                print "    Mounting %s" %self.mount_point
                cmd = "mount %s" %self.mount_point
                runcmd(cmd)
        return

    def write_config(self):
        print "    Writing /etc/storage.d conf file"
        self.config.add_section('storage')
        self.config.set('storage','uuid',self.uuid)
        self.config.set('storage','mountpoint',self.mount_point)
        self.config.set('storage','fstype',self.fstype)
        self.config.set('storage','shareable','True')
        self.config.set('storage','mmount',self.myth_mount)
        self.config.set('storage','storage_groups',self.dir_sg)
        self.config.set('storage','disk_num',self.disk_num)

        print "       %s" %self.conf_file
        with open(self.conf_file, 'wb') as self.conf_file:
            self.config.write(self.conf_file)
        return

def reconstruct_mounts(no_mount):
    print "\nRecreating disks from contents of /etc/storage.d/"
    for conf_file in glob.glob('%s/*.conf' %storage_dir):
        print "\n"
        cf = reconstruct_path(conf_file)

        #print cf.get_conf()
        #print cf.get_uuid()
        print "    Recreating %s" %cf.get_mount_point()
        #print cf.get_shareable()
        #print cf.get_is_myth_mount()
        #print cf.get_disk_num()

        cf.create_mount_point()
        cf.add_fstab()
        cf.symlink_disk()
        cf.mount_disk(no_mount)

    print "\n\nDone recreating disks.\n"
    pass


def usage():
    help='''
    add_storage.py finds and sets up disks for MythTV usage.
    It's a powerful tool that could destroy data if not used correctly,
        please be careful.

    Scanned disks are ignored if they are mounted or have been
        previously skipped by add_storage.py.

    The file system type for disks added by add_storage.py is
        automatically set to the type selected for the data partition
        at install.

    Normal operations without options include (in this order):
        Partition the disk
        Format the disk
        Add disk to /etc/fstab
        Mount the disk
        Create the directories
            (if user enables MythTV Storage Groups)
        Write out the disk config file to /etc/storage.d/
        Create disk# symlink at /data/storage/
        Create /myth symlink (if applicable)
        Create MythTV Storage Group paths in MythTV database
            (if user enables MythTV Storage Groups)

    Options:
    --add_sg:           Create the MythTV Storage Group directories and
                           database entries for database backups, TV
                           recordings, photos, music, streaming, videos
                           and artwork.
    -h, --help:         Show this help message.
    --new_init:         Erase the list of known disks and rescan.
    --no_destruction:   Will not partition or format the disk.
                           All other normal operations will be performed.
                           Can be used to import disks from other systems
                           however, add_storage.py only works with the first
                           partition on a disk and ignores all others.
    --no_mount:         Do not mount the disk.
                           All other normal operations will be performed.
    --reconstruct:      Recreate mount point, /myth symlink, fstab entry,
                           /data/storage/disk# symlink, and mount the disk.
                           --no_mount is the only option that works with
                           --reconstruct.
    --reconstruct_sg:   Recreate the MythTV Storage Group directories and
                           database entries if they don't exist.
                           No other options work with --reconstruct_sg.
    --report:           Scan disks and print new found disks.
    '''
    print help
    sys.exit(0)




if __name__ == "__main__":
    scan_only = False
    myth_mount = False
    no_mount = False
    destruction = True
    install_call = False
    dir_sg = False
    reconstruct = False
    reconstruct_sg = 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 or "-h" 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 "--add_sg" in sys.argv:
        dir_sg = True

       #there is no distinction between FE and BE sg anymore
       #but leaving these for backwards compatibility
    if "--add_fe_sg" in sys.argv:
        dir_sg = True

    if "--add_be_sg" in sys.argv:
        dir_sg = True

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

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

    if "--double_myth" in sys.argv:
        myth_main(no_mount, install_call, dir_sg)
    elif reconstruct == True:
        reconstruct_mounts(no_mount)
    elif reconstruct_sg == True:
        reconstruct_storagegroups()
    else:
        main(scan_only, destruction, no_mount, install_call, dir_sg)