Server Download File

[safesploit]

To clarify the logic is server.py runs download_file() function which then relies on the upload_file() function on client.py.

The following additional functionality has been implemented when server.py wants to download a file from the remote machine.

• File upload errors (client.py)
• Use context manager (utilise with statement; ensuring the file is closed properly, in the event of exceptions/errors)
• Implement file chunking: For large files, it's advisable to send the file in smaller chunks rather than reading and sending the entire file content at once. This allows for more efficient memory usage and prevents potential issues with large file transfers. Using a loop, you can modify the function to read and send the file in chunks until the entire file is uploaded.
• Progress updates

client.py

import hashlib
import os

def calculate_checksum(file_name):
    sha256_hash = hashlib.sha256()

    with open(file_name, 'rb') as file:
        for chunk in iter(lambda: file.read(4096), b''):
            sha256_hash.update(chunk)

    return sha256_hash.hexdigest()

def upload_file(target, file_name):
    try:
        file_size = os.path.getsize(file_name)
        chunk_size = 1024  # Set the chunk size for file transfer
        total_bytes = 0  # Track the total bytes transferred

        checksum = calculate_checksum(file_name)

        with open(file_name, 'rb') as file:
            while True:
                chunk = file.read(chunk_size)
                if not chunk:
                    break  # Exit the loop when the end of file is reached

                try:
                    target.sendall(chunk)
                    total_bytes += len(chunk)

                    # Calculate and display the upload progress
                    progress = (total_bytes / file_size) * 100
                    print(f"Uploading {file_name}: {progress:.2f}% completed")

                except Exception as e:
                    print(f"An error occurred during file upload: {str(e)}")
                    break

        print(f"File '{file_name}' uploaded successfully. Total bytes transferred: {total_bytes}")

        # Verify checksum on the server side
        reliable_send(target, checksum)
        verification_status = reliable_recv(target)
        if verification_status:
            print("Checksum verification successful.")
        else:
            print("Checksum verification failed.")

    except FileNotFoundError:
        print(f"Error: File '{file_name}' not found.")
    except IOError as e:
        print(f"Error reading file '{file_name}': {str(e)}")
    except Exception as e:
        print(f"An error occurred while handling file '{file_name}': {str(e)}")

server.py

import hashlib

def calculate_checksum(data):
    sha256_hash = hashlib.sha256()
    sha256_hash.update(data)
    return sha256_hash.hexdigest()

def download_file(target, file_name, type):
    try:
        f = open(file_name, 'wb')
        if type == 1:
            target.settimeout(2)
            chunk = target.recv(1024)
            while chunk:
                f.write(chunk)
                try:
                    chunk = target.recv(1024)
                except socket.timeout as e:
                    break
        elif type == 2:
            target.settimeout(5)
            try:
                chunk = target.recv(10485760)
            except:
                pass

            while chunk:
                f.write(chunk)
                try:
                    chunk = target.recv(10485760)
                except:
                    break
        target.settimeout(None)
        f.close()

        # Verify checksum after file download
        file_data = open(file_name, 'rb').read()
        received_checksum = reliable_recv(target)
        calculated_checksum = calculate_checksum(file_data)

        if received_checksum == calculated_checksum:
            reliable_send(target, True)  # Send checksum verification success status
            print(f"File '{file_name}' downloaded successfully. Checksum verification successful.")
        else:
            reliable_send(target, False)  # Send checksum verification failure status
            print(f"File '{file_name}' downloaded successfully. Checksum verification failed.")

    except Exception as e:
        print(f"An error occurred during file download: {str(e)}")