import datetime
import pexpect
import os
import logging
import sys
import tarfile
import os.path
import gzip
import psycopg2

class LogCollector:

    def __init__(self):
        pass
        self.logger = logging.getLogger()

    def get_connection(self):
        conn = psycopg2.connect(
            host="",
            database="",
            user="",
            password="")
        return conn

    def run_command_and_collect_output(self, site_name, site_ip, password, command):
        #print(site_name)
        #print(site_ip)
        #print(password)
        #print(command)
        #conn = self.get_connection()
        # create a cursor
        #try:
        #    cur = conn.cursor()
        
            # execute a statement
        #    print('PostgreSQL database version:')
        #    cur.execute('SELECT version()')

            # display the PostgreSQL database server version
        #    db_version = cur.fetchone()
        #    print(db_version)
       
            # close the communication with the PostgreSQL
        #    cur.close()
        #except (Exception, psycopg2.DatabaseError) as error:
        #    print(error)
        #finally:
        #    if conn is not None:
        #        conn.close()
        #        print('Database connection closed.')

        source_dir = site_name
        if not os.path.exists(source_dir):
            os.makedirs(source_dir)

        ssh_prefix = 'ssh -o StrictHostKeyChecking=no -t sysadmin@' + site_ip
        debugLog = site_name + "/cmd-output.log"
        fp = open(debugLog, "a")

        parts = command.split(" ")

        which_cmd = ssh_prefix + " which " + parts[0] 
        cmds = []
        cmds.append(which_cmd)
        all_lines = self.run_command_get_all_lines(cmds, password, self.logger)
        raw_cmd = ''
        for line in all_lines.split("\n"):
            line = line.rstrip().lstrip()
            if line and 'Connection' not in line:
                raw_cmd = line
                break

        #kubectl get ns --kubeconfig=/etc/kubernetes/admin.conf | grep wsbomagj | awk '{print $1}' | xargs kubectl delete ns --kubeconfig=/etc/kubernetes/admin.conf

        #cmd = ssh_prefix + " /usr/sbin/" + raw_cmd
        raw_cmd = raw_cmd + " " + ' '.join(parts[1:])
        if 'system' in raw_cmd:
            cmd = ssh_prefix + " source /etc/platform/openrc; " + raw_cmd
        elif 'kubectl' in raw_cmd:
            cmd = ssh_prefix + " " + raw_cmd + " --kubeconfig=/etc/kubernetes/admin.conf"
        else:
            cmd = ssh_prefix + " " + raw_cmd
        print(cmd)
        self._execute_cmd_and_save_op(cmd, raw_cmd, password, fp)
        fp.close()

    def collect_samsung_logs(self, site_name, site_ip, password):
        source_dir = site_name
        if not os.path.exists(source_dir):
            os.makedirs(source_dir)

        ssh_prefix = 'ssh -o StrictHostKeyChecking=no -t sysadmin@' + site_ip
        debugLog = site_name + "/DebugLog"
        fp = open(debugLog, "w")

        # 1. uptime
        raw_cmd = " uptime "
        uptime = ssh_prefix + raw_cmd
        self._execute_cmd_and_save_op(uptime, raw_cmd, password, fp)

        # 2. ls -l /var/lib/systemd/coredump
        raw_cmd = " ls -l /var/lib/systemd/coredump "
        coredump_ls = ssh_prefix + raw_cmd
        self._execute_cmd_and_save_op(coredump_ls, raw_cmd, password, fp)

        # 3. fm alarm-list
        raw_cmd = " fm alarm-list"
        fm_alarm_list = ssh_prefix + " source /etc/platform/openrc; " + raw_cmd
        self._execute_cmd_and_save_op(fm_alarm_list, raw_cmd, password, fp)

        # 4. date
        raw_cmd = " date  "
        date = ssh_prefix + raw_cmd
        self._execute_cmd_and_save_op(date, raw_cmd, password, fp)

        # 5. build.info
        raw_cmd = " cat /etc/build.info "
        build_info = ssh_prefix + raw_cmd
        self._execute_cmd_and_save_op(build_info, raw_cmd, password, fp)

        # 6. ip link
        raw_cmd = " /usr/sbin/ip link "
        ip_link = ssh_prefix + raw_cmd
        self._execute_cmd_and_save_op(ip_link, raw_cmd, password, fp)
        fp.close()

        # 7. pmond.log
        pmond_log = '/var/log/pmond.log'
        cmds = []
        cmd = 'scp -o StrictHostKeyChecking=no sysadmin@[' + site_ip + ']:' + pmond_log + " " + site_name + "/."
        cmds.append(cmd)
        self.run_commands_scp(cmds, password, self.logger)

        # 8. /var/log/kern.log
        kern_logs_cmd = ssh_prefix + " ls /var/log/kern.log*"
        cmds = [kern_logs_cmd]
        kern_logs = self.run_command_get_all_lines(cmds, password, self.logger)
        print(kern_logs)
        for kern_log in kern_logs.split('\n'):
            if 'Connection to' not in kern_log:
                kern_log = kern_log.rstrip().lstrip()
                kern_parts = kern_log.split("\t")
                print("Kern parts:" + str(len(kern_parts)))
                if len(kern_parts) == 1:
                    kern_parts = kern_log.split(" ")
                for p in kern_parts:
                    if p != '' and p != '\n' and p != '\t' and 'Connection' not in p:
                        cmds = []
                        print(p)
                        cmd = 'scp -o StrictHostKeyChecking=no sysadmin@[' + site_ip + ']:' + p + " " + site_name + "/."
                        print(cmd)
                        cmds.append(cmd)
                        self.run_commands_scp(cmds, password, self.logger)

        # 6. /var/log/user.log*
        user_logs_cmd = ssh_prefix + " ls /var/log/user.log*"
        cmds = [user_logs_cmd]
        user_logs = self.run_command_get_all_lines(cmds, password, self.logger)
        print(user_logs)
        for user_log in user_logs.split('\n'):
            if 'Connection to' not in user_log:
                cmds = []
                user_log = user_log.rstrip().lstrip()
                user_parts = user_log.split("\t")
                print("User parts:" + str(len(user_parts)))
                if len(user_parts) == 1:
                    user_parts = user_log.split(" ")
                for p in user_parts:
                    if p != '' and p != '\n' and p != '\t' and 'Connection' not in p:
                        print(p)
                        cmd = 'scp -o StrictHostKeyChecking=no sysadmin@[' + site_ip + ']:' + p + " " + site_name + "/."
                        print(cmd)
                        cmds.append(cmd)
                        self.run_commands_scp(cmds, password, self.logger)

        # 7. /var/log/pods/*adpf*
        pod_logs_cmd = ssh_prefix + " ls /var/log/pods/ | grep adpf"
        cmds = [pod_logs_cmd]
        pod_logs = self.run_command_get_all_lines(cmds, password, self.logger)
        print(pod_logs)
        for pod_log in pod_logs.split('\n'):
            cmds = []
            if ((pod_log is not None) and (pod_log != '')):
                cmd = 'scp -o StrictHostKeyChecking=no -r sysadmin@[' + site_ip + ']:/var/log/pods/' + pod_log + " " + site_name + "/."
                cmds.append(cmd)
                self.run_commands_scp(cmds, password, self.logger)

        # 8. /var/lib/systemd/coredump/*
        core_dump_cmd = ssh_prefix + " ls /var/lib/systemd/coredump "
        cmds = [core_dump_cmd]
        core_dump_logs = self.run_command_get_all_lines(cmds, password, self.logger)
        print(core_dump_logs)
        #for core_dump_log in core_dump_logs.split('\n'):
        cmds = []
        cmd = 'scp -o StrictHostKeyChecking=no -r sysadmin@[' + site_ip + ']:/var/lib/systemd/coredump/' + " " + site_name + "/."
        cmds.append(cmd)
        self.run_commands_scp(cmds, password, self.logger)

        # create tar ball
        date = datetime.datetime.now().strftime("%b-%d-%Y")
        tar_file_name = site_name + '-' + date + '.tar'
        gzip_file_name = tar_file_name + '.gz'
        os.system('tar -cvf ' + tar_file_name + ' ./' + site_name)
        os.system('gzip ' + tar_file_name)
        print(gzip_file_name + " created.")
        print("Upload it to Samsung FTP site and send a note to Samsung team. FTP coordinates:")
        #print("    url: https://zone.samsungtelecom.com/")
        #print("    username:vzw_femto_ndbm_jumpbox")
        #print("    password:S@msung123")
        suser = "verizon_snap"
        shost = "zone.samsungtelecom.com"
        spasswd = "VZWSn@p786"
        print("     ur: "+shost+"/ 65.169.250.6")
        print("     username:"+suser)
        print("     password:"+spasswd)
        # sftp the file up to samsung
        remotepath = '/VZW_4G_5G/Controller\ Logs'
        p = pexpect.spawn('sftp %s@%s' %(suser,shost))
        p.logfile = sys.stdout.buffer
        try:
            p.expect('(?i)password:')
            x = p.sendline(spasswd)
            x = p.expect(['Permission denied','sftp>'])
            if x == 0:
                print ('Permission denied for password:')
                print (spasswd)
                p.kill(0)
            else:
                x = p.sendline('cd ' + remotepath)
                x = p.expect('sftp>')
                x = p.sendline('put ' + gzip_file_name)
                x = p.expect('sftp>')
                x = p.isalive()
                x = p.close()
                retval = p.exitstatus
        except EOF:
            print (str(p))
            print ('SFTP file transfer failed due to premature end of file.')
        except TIMEOUT:
            print (str(p))
            print ('SFTP file transfer failed due to timeout.')

    def _handle_collect(self, site_name, site_ip, password):
        ssh_prefix = 'ssh -o StrictHostKeyChecking=no -t sysadmin@' + site_ip
        collect_cmd = ssh_prefix + " collect all"
        cmds = []
        cmds.append(collect_cmd)
        all_lines = self._run_collect(cmds, password)
        for line in all_lines.split("\n"):
            if 'creating' in line:
                parts = line.split(" ")
                collect_file_location = parts[3]
                print("Collect file location:" + collect_file_location)

                cmds = []
                cmd = 'scp -o StrictHostKeyChecking=no sysadmin@[' + site_ip + ']:' + collect_file_location + " " + site_name + "/."
                cmds.append(cmd)
                self.run_commands_scp(cmds, password, self.logger)

    def collect_windriver_logs_steps(self, site_name, site_ip, password):
        source_dir = site_name
        if not os.path.exists(source_dir):
            os.makedirs(source_dir)

        # 1. collect
        self._handle_collect(site_name, site_ip, password)

        # 2. network interface details
        ssh_prefix = 'ssh -o StrictHostKeyChecking=no -t sysadmin@' + site_ip
        port_list_cmd = ssh_prefix + " \"source /etc/platform/openrc; system host-port-list controller-0 | grep -v name | grep -v '+' | cut -d ' ' -f 4\""
        cmds = []
        cmds.append(port_list_cmd)
        all_lines = self.run_command_get_all_lines(cmds, password, self.logger)

        for interface in all_lines.split("\n"):
            if 'Connection' not in interface:
                print("Interface..." + interface)
                interface_file = site_name + "/" + interface + ".txt"
                fp = open(interface_file, "a")

                raw_cmd = " /usr/sbin/ethtool -i " + interface
                cmd1 = ssh_prefix + raw_cmd
                print(cmd1)
                self._execute_cmd_and_save_op(cmd1, raw_cmd, password, fp)

                raw_cmd = " /usr/sbin/ethtool -l " + interface
                cmd1 = ssh_prefix + raw_cmd
                self._execute_cmd_and_save_op(cmd1, raw_cmd, password, fp)

                raw_cmd = " /usr/sbin/ethtool -S " + interface
                cmd1 = ssh_prefix + raw_cmd
                self._execute_cmd_and_save_op(cmd1, raw_cmd, password, fp)

                fp.close()

        # 3. interrupts
        interrupt_file = site_name + "/interrupts.txt"
        fp = open(interrupt_file, "a")
        raw_cmd = " cat /proc/interrupts"
        cmd1 = ssh_prefix + raw_cmd
        self._execute_cmd_and_save_op(cmd1, raw_cmd, password, fp)
        fp.close()

        debugLog = site_name + "/wr-debug.log"
        fp = open(debugLog, "w")

        # 4. date
        raw_cmd = " date  "
        date = ssh_prefix + raw_cmd
        self._execute_cmd_and_save_op(date, raw_cmd, password, fp)

        # 5. ipmitool
        raw_cmd = " sudo ipmitool sel time get "
        cmd1 = ssh_prefix + raw_cmd
        print(cmd1)
        self._execute_cmd_and_save_op(cmd1, raw_cmd, password, fp, additional_password=True)

        # 6. uptime
        raw_cmd = " uptime "
        uptime = ssh_prefix + raw_cmd
        self._execute_cmd_and_save_op(uptime, raw_cmd, password, fp)

        # 7. host_if_list
        raw_cmd = "system host-if-list controller-0"
        host_if_list = ssh_prefix + " source /etc/platform/openrc; " + raw_cmd
        self._execute_cmd_and_save_op(host_if_list, raw_cmd, password, fp)

        # 8. interface_network_list
        raw_cmd = "system interface-network-list controller-0"
        interface_network_list = ssh_prefix + " source /etc/platform/openrc; " + raw_cmd
        self._execute_cmd_and_save_op(interface_network_list, raw_cmd, password, fp)

        # 9. fm alarm-list
        raw_cmd = " fm alarm-list"
        fm_alarm_list = ssh_prefix + " source /etc/platform/openrc; " + raw_cmd
        self._execute_cmd_and_save_op(fm_alarm_list, raw_cmd, password, fp)

        # 5. fm event-list
        #raw_cmd = " fm event-list"
        #fm_event_list = ssh_prefix + " source /etc/platform/openrc; " + raw_cmd
        #self._execute_cmd_and_save_op(fm_event_list, raw_cmd, password, fp)

        # 10. sw-patch query
        raw_cmd = " /usr/sbin/sw-patch query"
        sw_patch_query = ssh_prefix + " source /etc/platform/openrc; " + raw_cmd
        self._execute_cmd_and_save_op(sw_patch_query, raw_cmd, password, fp)
        fp.close()

        # 11. drive details
        drive_list_cmd = ssh_prefix + " \"source /etc/platform/openrc; system host-disk-list controller-0 | grep \"/dev/disk\" | awk '{print $4}' \""
        cmds = []
        cmds.append(drive_list_cmd)
        all_lines = self.run_command_get_all_lines(cmds, password, self.logger)

        if not os.path.exists(site_name + "/dev"):
            os.makedirs(site_name + "/dev")

        for drive in all_lines.split("\n"):
            if 'Connection' not in drive:
                print("Drive..." + drive)
                drive_file = site_name + drive + ".txt"
                fp = open(drive_file, "w")
                raw_cmd = " sudo smartctl -a  " + drive
                cmd1 = ssh_prefix + raw_cmd
                print(cmd1)
                self._execute_cmd_and_save_op(cmd1, raw_cmd, password, fp, additional_password=True)
                fp.close()

        # create tar ball
        tar_file_name = source_dir + '.tar'
        gzip_file_name = tar_file_name + '.gz'
        os.system('tar -cvf ' + tar_file_name + ' ./' + source_dir)
        os.system('gzip ' + tar_file_name)
        #with tarfile.open(tar_file_name, "w:gz") as tar:
        #    tar.add(source_dir, arcname=os.path.basename(source_dir))
        print(gzip_file_name + " created.")

    def collect_windriver_logs(self, site_name, site_ip, password):
        source_dir = site_name
        if not os.path.exists(source_dir):
            os.makedirs(source_dir)

        ssh_prefix = 'ssh -o StrictHostKeyChecking=no -t sysadmin@' + site_ip

        # Running collect
        cmds = []
        cmd = ssh_prefix + " collect"
        cmds.append(cmd)
        all_lines = self._run_collect(cmds, password)
        for line in all_lines.split("\n"):
            if 'creating' in line:
                parts = line.split(" ")
                collect_file_location = parts[3]
                print("Collect file location:" + collect_file_location)

    def _execute_cmd_and_save_op(self, cmd, raw_cmd, password, fp, additional_password=False):
        cmds = []
        #print(cmd)
        cmds.append(cmd)
        all_lines = self.run_command_get_all_lines(cmds, password, self.logger, additional_password=additional_password)
        #print(all_lines)
        fp.write(raw_cmd + ":\n")
        fp.write("-------------------------\n")
        for line in all_lines.split("\n"):
            if 'Connection' not in line:
                fp.write(line)
                fp.write("\n")
        fp.write("\n")

    def _run_collect(self, commands, host_password, cmd_timeout=None):
        all_lines = []
        for command in commands:
            print(command)
            child = pexpect.spawn(command)
            child.timeout=cmd_timeout
            try:
                child.expect_exact(['password: '], timeout=cmd_timeout)
                child.sendline(host_password)
                child.expect_exact(['[sudo] password for sysadmin:'], timeout=cmd_timeout)
                child.sendline(host_password)
                #child.sendline("\n")
                line = child.readline()
                print(line)
                line = child.readline()
                print(line)
                all_lines = child.read()
                all_lines = all_lines.rstrip().lstrip()
                all_lines = all_lines.decode('utf-8').replace('\r\n', '\n')
                print(all_lines)
            except:
                print(str(child))
        return all_lines

    def run_command_get_all_lines(self, commands, host_password, logger, block=False, timeout=None, additional_password=False):
        #print("Inside _run_command_get_all_lines")
        all_lines = []
        for command in commands:
            print("Executing.." + command)
            child = pexpect.spawn(command)
            child.timeout=timeout
            try:
                i = child.expect(['password: ','Connection refused\r\r\n'], timeout=timeout)
                if i == 0:
                    child.sendline(host_password)
                    if additional_password:
                        child.expect_exact("Password:")
                        child.sendline(host_password)
                    all_lines = child.read()
                    all_lines = all_lines.rstrip().lstrip()
                    all_lines = all_lines.decode('utf-8').replace('\r\n', '\n')
                    #logger.info(all_lines)
                if i == 1:
                    print("Connection refused")
            except:
                print(str(child))
        return all_lines

    def run_commands_scp(self, commands, host_password, logger, block=False, timeout=None):
        successful = False
        #print("Inside run_command_scp")
        #print(commands)
        for command in commands:
            print("Executing.." + command)
            try:
                child = pexpect.spawn(str(command))
            except:
                print(str(child))
            child.timeout = None
            try:
                i = child.expect([b'', 'Are you .*(yes/no)? ', pexpect.EOF], timeout=None)
                if i == 0:
                    child.sendline('') 
                if i == 1:
                    child.sendline('yes')
                child.expect(['password: '], timeout=timeout)
                child.sendline(host_password)
                all_lines = child.read()
                all_lines = all_lines.rstrip().lstrip()
                all_lines = all_lines.decode('utf-8').replace('\r\n', '\n')
                successful = True # Tentative
                for line in all_lines.split("\n"):
                    #print(line)
                    if re.search('error', line, re.IGNORECASE):
                        successful = False 
                    if re.search('unable', line, re.IGNORECASE):
                        successful = False 
            except:
            	pass
                #print(str(child))
        #print("Status:" + str(successful))
        return successful

if __name__ == '__main__':

    if len(sys.argv) < 5:
        print("Samsung logs:")
        print("  python collect-logs.py samsung <site-name> <site-ip> <password>")
        print("WindRiver logs:")
        print("  python collect-logs.py windriver <site-name> <site-ip> <password>")
        print("One-off command o/p:")
        print("  python collect-logs.py command <site-name> <site-ip> <password> \"<command>\"")
        print("  python3 collect-logs.py command longmeadow-oct8 2001:4888:2a10:e0ad:122:40a:0:f400 password \"ifconfig\"")
        print("  python3 collect-logs.py command longmeadow-oct8 2001:4888:2a10:e0ad:122:40a:0:f400 password \"ip link\"")
        exit(0)

    logs_for = sys.argv[1]
    site_name = sys.argv[2]
    site_ip = sys.argv[3]
    password = sys.argv[4]

    print("Option chosen:" + logs_for)
    print("Site:" + site_name)
    print("Site IP:" + site_ip)
    print("Site Password:" + password)
    log_collector = LogCollector()

    if logs_for == 'samsung':
        log_collector.collect_samsung_logs(site_name, site_ip, password)
    elif logs_for == 'windriver':
        log_collector.collect_windriver_logs_steps(site_name, site_ip, password)
    else:
        command = sys.argv[5]
        #print(command)
        log_collector.run_command_and_collect_output(site_name, site_ip, password, command)