client.py

__author__ = "Akshay Joshi"

import os
import time

# Set root directory of the repo into PATH
import sys
sys.path.append('../Question_3/')

# Import python socket library & custom file to track static variables
from socket import *
from parameters import *


class Client:
    """
        Code to simulate operations of a Client.
        Operations implemented: 
            1. Connection Setup
            2. Upload
            3. Download
            4. Hash Generation & Matching
    """
    
    def __init__(self):
        self.server_port = PORT  # Server Port
        self.server_url = URL    # Server URL
        self.commands = COMMANDS    # File Commands
        self.seperator = SEP  # Separator token for file name & file size strings
        self.buffer_size = BUFFER_CAPACITY  # Buffer size
        self.file_name = "hacker_cat.jpg"
        self.file_size = os.path.getsize(self.file_name)
        self.received_files_path = RECEIVED_FILES_PATH


    def match_md5_hash(self):
        # Receive MD5 hash generated by the server & compares it against the local MD5 hash
        received = self.clientSocket.recv(self.buffer_size).decode()

        print(f"\nClient MD5 (Original Image): {self.hacker_cat_img_md5}\n")
        print(f"\nServer MD5 (Uploaded/Downloaded Image): {received}\n")

        if received == self.hacker_cat_img_md5:
            print("\n-----Success-----\n")

        else:
            print("\n-----Fail-----\n")


    def list_files(self):
        self.clientSocket.send(COMMANDS[0].encode(FORMAT))

        msg = None
        availble_filenames = []
        while msg != 'DONE':
            msg = self.clientSocket.recv(self.buffer_size).decode()
            availble_filenames.append(msg)

        for filenames in availble_filenames[:-1]:
            print(f"{filenames}\n")
        

    def upload(self):
        self.clientSocket.send(COMMANDS[1].encode(FORMAT))
        msg = self.clientSocket.recv(self.buffer_size).decode()
        print(msg)

        # Send filename & filesize to the server
        print("\nSending filename & filesize to the server...")
        print(f"Filename: {self.file_name}, Filesize: {self.file_size}\n")
        self.clientSocket.send(f"{self.file_name}{self.seperator}{self.file_size}".encode(FORMAT))

        msg = self.clientSocket.recv(self.buffer_size).decode()
        print(msg)

        # Open the desired image to upload
        with open(self.file_name, 'rb') as hacker_cat_img:

            # Send the data
            while True:
                data_bytes = hacker_cat_img.read(self.buffer_size)
                md5().update(data_bytes)

                if not data_bytes:
                    break   # All the data sent or empty file

                # Using sendall instead of send as this method continues to send data 
                # from data_bytes until either all data has been sent or an error occurs.

                # Don't need to use encode() as we are already loading the file in 'rb'
                # mode
                self.clientSocket.sendall(data_bytes)

            # Adding a small time delay
            time.sleep(0.1)

            # Sending an End of File flag to notify the receiver that the 
            # stream completes here.
            self.clientSocket.send('EOF'.encode()) 
            
            """
            # Compute MD5 hash of the image to be sent - deprecated
            for block in iter(lambda: hacker_cat_img.read(self.buffer_size), b""):
                md5().update(block)
            """

            # Returns the hex string representation for the digest
            self.hacker_cat_img_md5 = md5().hexdigest()  
            hacker_cat_img.close()


    def download(self):
        self.clientSocket.send(COMMANDS[2].encode(FORMAT))
        msg = self.clientSocket.recv(self.buffer_size).decode()
        print(msg)

        # Send file ID to the server
        print("\nSending file ID to the server...\n")
        self.clientSocket.send(self.hacker_cat_img_md5.encode(FORMAT))

        print(f"File ID that is sent to server is: {self.hacker_cat_img_md5}")

        try:
            hacker_cat_temp_file = os.path.join(self.received_files_path, 'downloaded_image_from_server') + '.jpg'
            with open(hacker_cat_temp_file, 'wb') as hacker_cat_img:
                while True:
                    data_bytes =  self.clientSocket.recv(self.buffer_size)
                    md5().update(data_bytes)

                    # When the End of File flag is reached (which means file transfer is complete)
                    if data_bytes == 'EOF'.encode():
                        break   

                    hacker_cat_img.write(data_bytes)
                self.hacker_cat_img_md5 = md5().hexdigest()
                hacker_cat_img.close()

        except Exception as e:
            print(e)


    def communicate_to_server(self):
        # Create TCP socket for future connections
        self.clientSocket = socket(AF_INET, SOCK_STREAM)

        try:
            # Connect the client to the specified server
            self.clientSocket.connect((self.server_url, self.server_port))
            print("Client connected to server: " + self.server_url + " : " + str(self.server_port))

        except:
            print("\nThe server is not responding to requests!")
            print("\nPlease start the server first & try again")
            exit(99)

        # Send the 'LIST_FILES' command
        print("\nSending 'LIST_FILES' command to the server...\n")
        self.list_files()

        # Send the 'UPLOAD' command
        print("\nSending 'UPLOAD' command to the server...\n")
        self.upload()

        # Match MD5 hashes from client & server
        print("\nMatching MD5 hashes (Client & Server) after UPLOAD...\n")
        self.match_md5_hash()

        # List the updated files
        print("\nSending 'LIST_FILES' command to the server once again...\n")
        self.list_files()

        # Download a specific file match by it's file ID (first 32 chars in the filename)
        print("\nSending 'DOWNLOAD' command to the server...\n")
        self.download()

        # Match MD5 hashes from client & server
        print("\nMatching MD5 hashes (Client & Server) after DOWNLOAD...\n")
        self.match_md5_hash()

        # Close TCP connection
        print("\nClosing Client Socket...\n")
        self.clientSocket.close()

# Driver code
if __name__ == "__main__":
    client = Client()
    client.communicate_to_server()