#!/usr/bin/python
#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.
#
# Version 2.1.0

import dbus
import pickle
import subprocess
import sys,os,re
import random,string
import configparser
from configparser import ConfigParser
import glob
import logging
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 list(SG_MAP.keys()):
    FS_LIST.append(SG_MAP[key])

class disk_device:
    def  __init__(self,device,storage_dir):
        block_dev = bus.get_object("org.freedesktop.UDisks2", device)

        self.block_path = block_dev.Get('org.freedesktop.UDisks2.Block', 'Device', dbus_interface='org.freedesktop.DBus.Properties')
        self.block_path = bytearray(self.block_path).replace(b'\x00', b'').decode('utf-8')
        logging.info("Device: %s", self.block_path)
        self.read_only = self.get_is_readonly(block_dev)
        logging.info("ReadOnly: %s", self.read_only)
        self.device_file_path = self.get_device_file_path(block_dev)
        logging.info("Device File Path: %s", self.device_file_path)
        self.device_id = block_dev.Get('org.freedesktop.UDisks2.Block', 'Id', dbus_interface='org.freedesktop.DBus.Properties')
        logging.info("Device Id: %s", self.device_id)
        self.is_device = self.get_is_device()
        logging.info("Is Device: %s", self.is_device)

        self.drive = block_dev.Get('org.freedesktop.UDisks2.Block', 'Drive', dbus_interface='org.freedesktop.DBus.Properties')
        logging.info("Drive: %s", self.drive)
        drive_dev = bus.get_object("org.freedesktop.UDisks2", self.drive)

        self.storage_dir = storage_dir
        self.top_mount_dir = "/data/storage"
        self.config = configparser.RawConfigParser()
        self.fs_map = self.get_fsmap()

        self.vendor = drive_dev.Get('org.freedesktop.UDisks2.Drive', 'Vendor', dbus_interface='org.freedesktop.DBus.Properties')
        logging.info("Vendor: %s", self.vendor)
        self.model = drive_dev.Get('org.freedesktop.UDisks2.Drive', 'Model', dbus_interface='org.freedesktop.DBus.Properties')
        logging.info("Model: %s", self.model)
        self.device_size = drive_dev.Get('org.freedesktop.UDisks2.Drive', 'Size', dbus_interface='org.freedesktop.DBus.Properties')
        logging.info("Drive Size: %s", self.device_size)
        self.serial_number = self.get_serial_number(drive_dev)
        logging.info("Serial: %s", self.serial_number)
        self.is_optical = self.get_is_optical_disc(drive_dev)
        logging.info("Is Optical: %s", self.is_optical)

        self.mmount = False
        self.dir_sg = False

        try:
            self.f = block_dev.Get('org.freedesktop.UDisks2.Filesystem', 'MountPoints', dbus_interface='org.freedesktop.DBus.Properties')
            self.is_mounted = True
            self.f[0] = bytearray(self.f[0]).replace(b'\x00', b'').decode('utf-8')
            logging.info("MountPoints: %s", self.f[0])
        except:
            self.f = ['']
            self.is_mounted = False
            logging.info("MountPoints: %s", self.f[0])
        logging.info("Is Mounted: %s", self.is_mounted)

        try:
            self.partition_size = block_dev.Get('org.freedesktop.UDisks2.Partition', 'Size', dbus_interface='org.freedesktop.DBus.Properties')
        except:
            self.partition_size = 0
        logging.info("Partition Size: %s", self.partition_size)

        self.set_partition("1")
        logging.info("Block Partition: %s", self.block_partition)

        self.in_use = self.get_in_use()
        logging.info("In Use: %s", self.in_use)
        self.uuid=''
        self.new_mount_point=''
        self.disk_num=''

    def set_partition(self,partition):
        if self.is_device:
            if 'nvme' in self.block_path:
                self.block_partition = "%sp%s" %(self.block_path,partition)
            else:
                self.block_partition = "%s%s" %(self.block_path,partition)
        else:
            self.block_partition = self.block_path

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

    def get_is_readonly(self,block_dev):
        readonly = block_dev.Get('org.freedesktop.UDisks2.Block', 'ReadOnly', dbus_interface='org.freedesktop.DBus.Properties')
        if readonly == 0:
            return False
        else:
            return True

    def get_is_optical_disc(self,drive_dev):
        optical = drive_dev.Get('org.freedesktop.UDisks2.Drive', 'Optical', dbus_interface='org.freedesktop.DBus.Properties')
        if optical == 0:
            return False
        else:
            return True

    def get_is_device(self):
        #match = re.search(r'part\d+$', self.device_file_path)
        match = re.search(r'part\d+$', self.device_id)
        if match is None:
            return True
        else:
            return False

    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_serial_number(self,drive_dev):
        serial_number = drive_dev.Get('org.freedesktop.UDisks2.Drive', 'Serial', dbus_interface='org.freedesktop.DBus.Properties')
        random_string = os.urandom(5)
        if serial_number == '':
            serial_number = "".join( [random.choice(string.ascii_letters) for i in range(6)] )
            serial_number = "NSW%s" %serial_number
        return serial_number

    def get_device_file_path(self,block_dev):
        paths = block_dev.Get('org.freedesktop.UDisks2.Block', 'Symlinks', dbus_interface='org.freedesktop.DBus.Properties')
        try:
            for path in paths:
                path = bytearray(path).replace(b'\x00', b'').decode('utf-8')
                if path.startswith('/dev/disk/by-path'):
                    return path
        except:
            path = "None"

    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):
        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

        loop = True
        prompt = '''
        Format disk with what filesystem (xfs/ext4)?:'''

        while loop:
            str1 = input(prompt)

            if str1 in ['xfs','ext4']:
                loop = False
                fstype=str1
                break
        print("\n")

        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):
        #new_fstab_list=['UUID=', 'mount_point', 'auto', 'defaults', '0', '1']
        new_fstab_list=['UUID=', 'mount_point', self.new_fstype, '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(' ',''))

        #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]
        new_options = "nofail,x-systemd.device-timeout=10"

        #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
        if self.new_fstype == "xfs":
            new_fstab_list[5]="0"

        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')
        for i in fstab:
            f.write(i)
        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.keys():
            #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, 'w') 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)


#end of class


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




def scan_system():
    ud_manager_obj = bus.get_object("org.freedesktop.UDisks2", "/org/freedesktop/UDisks2")
    ud_manager = dbus.Interface(ud_manager_obj, 'org.freedesktop.DBus.ObjectManager')
    current_drive_list=[]
    for dev in ud_manager.GetManagedObjects().items():
        if dev[0].startswith("/org/freedesktop/UDisks2/block_devices"):
            logging.info(dev[0])
            drive = disk_device(dev[0],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.partition_size,system_drive.device_file_path)
            y_drive_hash = "%s_%s_%s" %(y.model,y.partition_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 = 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 = 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 = 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 = ConfigParser()
    num_list = []
    for conf_file in glob.glob('%s/*.conf' %storage_dir):
        parser.read(conf_file)
        try:
            disk_num = parser.get('storage', 'disk_num')
        except:
            print("\nSkipping " + conf_file + "is missing disk_num.")
            continue
        num_list.append(int(disk_num))
        num_list.sort()
    try:
        return num_list[-1]
    except:
        # conf file or disk_num is missing so fallback to /data/storage/disk# links
        for disk_name in glob.glob('/data/storage/disk*'):
            disk_num = disk_name.strip('/data/storage/disk')
            num_list.append(int(disk_num))
            num_list.sort()
        try:
            return num_list[-1]
        except:
            print("Couldn't find last disk number.")
            sys.exit(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)


    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 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:
                    DB = MythDB()
                    i.add_sg(DB,host,SG_MAP)

                print("-----")

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 = ConfigParser()
        parser.read(conf_file)
        try:
            mount_point = parser.get('storage', 'mountpoint')
        except:
            print("\nSkipping: " + conf_file + " is missing mountpoint")
            continue
        mmount = parser.getboolean('storage', 'mmount')
        try:
            removed = parser.getboolean('storage', 'removed')
        except:
            removed = False
        if removed:
            print("Skipping: " + mount_point + " - removed")
            continue
        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 as 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.keys():
                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] == '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 = ConfigParser()
        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.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()
        try:
            self.removed = parser.get('storage', 'removed')
        except:
            self.removed = False

    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 get_removed(self):
        return self.removed

    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.write(new_fstab_line)
        f.close()

    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)
        if os.path.islink(disk_ln):
            print("        Symlink %s exists. Skipping." %disk_ln)
        else:
            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.check_in_fstab(fstab,self.uuid) == True:
            print("    Found UUID of disk in fstab, will not add it")
        else:
            print("    Adding storage to fstab")
            #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_options = "nofail,x-systemd.device-timeout=10"
            new_fstab_list[0]="UUID=%s" %self.uuid
            new_fstab_list[1]=self.mount_point
            new_fstab_list[3]=new_options
            if self.fstype == "xfs":
                new_fstab_list[5]="0"
            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, 'w') 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)

        # skip if the disk was removed
        if cf.get_removed():
            continue
        #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 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, 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")

    logging.basicConfig(filename='/var/log/add_storage.log', filemode='w',
                        format='%(asctime)s - %(levelname)s - %(message)s',
                        datefmt='%y-%m-%d %H:%M:%S')
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    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 reconstruct == True:
        reconstruct_mounts(no_mount)
    elif reconstruct_sg == True:
        reconstruct_storagegroups()
    else:
        main(scan_only, destruction, no_mount, install_call, dir_sg)
