import subprocess
import os
from ftplib import FTP
import time

# Configurations
FTPorSFTPserver = 1  # 0 for FTP, 1 for SFTP
hostname = "c1report.comodo.com"
username = "c1report"
password = "paT7rObeseW0duLl"
dirpath = "reports"



timestamp = time.strftime("%Y%m%d_%H%M%S")  # Format: YYYYMMDD_HHMMSS


def check_directory():
    if os.path.exists("/Library/Application Support/COMODO/CCS"):
        return "CCSM"
    elif os.path.exists("/Library/Application Support/Xcitium Client - Security"):
        return "XCSM"
    else:
        return "EM"

log_type = check_directory()
print(log_type)
devicename = os.popen("hostname").read().split("-")[0]
filename = f"{devicename}_{timestamp}_{log_type}logs.zip"

# FTP upload function
def uploadingfileFTP(hostname, username, password, dirpath, filepath, filename):
    ftp = FTP()
    ftp.connect(hostname)
    ftp.login(username, password)
    ftp.cwd(dirpath)
    with open(filepath, 'rb') as file:
        ftp.storbinary(f"STOR {filename}", file)
    ftp.quit()
    print("Successfully uploaded the file to the FTP server")

# SFTP upload function using paramiko
def upload_sftp(hostname, username, password, dirpath, filepath, filename):
    transport = paramiko.Transport((hostname, 22))
    transport.connect(username=username, password=password)
    sftp = paramiko.SFTPClient.from_transport(transport)
    sftp.chdir(dirpath)
    sftp.put(filepath, filename)
    sftp.close()
    transport.close()
    print("Successfully uploaded the file to the SFTP server")

# Bash script to collect logs
bash_script = f"""#!/bin/bash
mkdir -p ~/Documents/{log_type}logs
mkdir -p ~/Documents/{log_type}logs/Logs
mkdir -p ~/Documents/{log_type}logs/SystemLogs
mkdir -p ~/Documents/{log_type}logs/DiagnosticReports

# List directories
ls /Library/LaunchDaemons/  > ~/Documents/{log_type}logs/FileList.txt
ls /Library/LaunchAgents/  >> ~/Documents/{log_type}logs/FileList.txt

# Copy logs
cp -a /var/log/ ~/Documents/{log_type}logs/SystemLogs/
cp -a /Library/Logs/DiagnosticReports/ ~/Documents/{log_type}logs/DiagnosticReports/

if [ {log_type} = "CCSM" ]; then
    ls /Applications/Comodo/CCS/ >> ~/Documents/{log_type}logs/FileList.txt
    ls /Library/Application Support/Comodo/CCS  >> ~/Documents/{log_type}logs/FileList.txt
    cp -a /Library/Application Support/COMODO/CCS/Logs/ ~/Documents/{log_type}logs/Logs/
    cp -a /Library/Application Support/Comodo/CCS/settings.plist ~/Documents/{log_type}logs/Logs/

elif [ {log_type} = "XCSM" ]; then
    ls "/Applications/Xcitium Client - Security.app/Contents/" >> ~/Documents/{log_type}logs/FileList.txt
    ls "/Library/Application Support/Xcitium Client - Security/"  >> ~/Documents/{log_type}logs/FileList.txt
    cp -a "/Library/Application Support/Xcitium Client - Security/Logs/" ~/Documents/{log_type}logs/Logs/
    cp -a "/Library/Application Support/Xcitium Client - Security/settings.plist" ~/Documents/{log_type}logs/Logs/
fi

zip -r /Library/Caches/{filename} ~/Documents/{log_type}logs -x "*.DS_Store"
rm -r ~/Documents/{log_type}logs
"""

# Ensure hostname is correctly parsed
if hostname.startswith("sftp://"):
    hostname = hostname.split("//")[1]
elif hostname.startswith("ftp://"):
    hostname = hostname.split("//")[1]
# If no protocol prefix exists, use the hostname as is
else:
    print("Warning: No protocol specified, using the hostname as is.")

# Now you can safely use the updated hostname in the sftpCommands
sftpCommands = f"""spawn sftp {username}@{hostname}
expect "password:"
send "{password}\\n"
expect "sftp>"
send "cd {dirpath}\\n"
expect "sftp>"
send "put /Library/Caches/{filename}\\n"
expect "sftp>"
send "exit\\n"
interact
"""


# Function to execute commands
def execute(arguments):
    p = subprocess.Popen(arguments, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    output, error = p.communicate()
    if output:
        print(output)

# Write scripts to disk
script_path = f"/Library/Caches/Collect{log_type}Logs.sh"
command_path = "/Library/Caches/sftpcommands.sh"

with open(script_path, 'w') as f:
    f.write(bash_script)

with open(command_path, 'w') as f:
    f.write(sftpCommands)

# Execute bash script to collect logs
execute(f"sh {script_path}")

# Filepath of the collected logs
filepath = f"/Library/Caches/{filename}"

# Check if file exists before uploading
if os.path.exists(filepath):
    if FTPorSFTPserver == 0:
        uploadingfileFTP(hostname, username, password, dirpath, filepath, filename)
    else:
        subprocess.call(["ssh", "-o", "StrictHostKeyChecking=no", f"{username}@{hostname}"])
        subprocess.call(["expect", command_path])

    # Clean up temporary files
else:
    print(f"File does not exist: {filepath}")
os.remove(script_path)
os.remove(filepath)
os.remove(command_path)
print(f"Logs are collected and uploaded successfully. File name on FTP(SFTP) is {filename}")