From 5a81257766f457603c5140f69d0731ed1733f973 Mon Sep 17 00:00:00 2001
From: Satyajit <jenasatyajit.sj@gmail.com>
Date: Thu, 2 Jan 2025 20:44:44 +0530
Subject: [PATCH] added input validation

---
 main.py | 146 ++++++++++++++++++++++++++++++++++++++++++++------------
 1 file changed, 116 insertions(+), 30 deletions(-)

diff --git a/main.py b/main.py
index 6710941..9fea6cd 100644
--- a/main.py
+++ b/main.py
@@ -14,6 +14,64 @@
 from bitarray import bitarray
 import requests
 import json
+import os
+
+class InputValidator:
+    @staticmethod
+    def validate_image_path(path: str, must_exist: bool = True) -> bool:
+        """Validates if the path is a valid image file"""
+        if not path:
+            raise ValueError("Path cannot be empty")
+            
+        valid_extensions = ('.png', '.jpg', '.jpeg', '.bmp')
+        if not path.lower().endswith(valid_extensions):
+            raise ValueError(f"File must be one of: {', '.join(valid_extensions)}")
+            
+        if must_exist and not os.path.exists(path):
+            raise FileNotFoundError(f"File does not exist: {path}")
+            
+        return True
+    
+    @staticmethod
+    def validate_message(message: str) -> bool:
+        """Validates the input message"""
+        if not message:
+            raise ValueError("Message cannot be empty")
+        
+        if len(message.encode('utf-8')) > 1000:  # size limit
+            raise ValueError("Message is too long (max 1000 bytes)")
+            
+        return True
+    
+    @staticmethod
+    def validate_destination_path(path: str) -> bool:
+        """Validates the destination path"""
+        try:
+            # Check if directory exists
+            directory = os.path.dirname(path)
+            if directory and not os.path.exists(directory):
+                raise ValueError(f"Directory does not exist: {directory}")
+                
+            # Check if path is writable
+            if os.path.exists(path):
+                if not os.access(path, os.W_OK):
+                    raise PermissionError(f"No write permission for: {path}")
+                    
+            return True
+        except Exception as e:
+            raise ValueError(f"Invalid destination path: {str(e)}")
+
+def get_validated_input(prompt_text: str, validator_func, error_msg: str = None):
+    """Generic function to get validated input"""
+    while True:
+        try:
+            user_input = input(prompt_text).strip()
+            validator_func(user_input)
+            return user_input
+        except (ValueError, FileNotFoundError, PermissionError) as e:
+            print(f"Error: {str(e)}")
+            if error_msg:
+                print(error_msg)
 
 def Encoder(Source, Message, Destination): 
     img = cv2.imread(Source)  
@@ -77,6 +135,10 @@ def send_metadata(metadata, token):
 
 def Decoder(Source):
     img = cv2.imread(Source)
+    if img is None:
+        print("Error: Image not found")
+        return
+    
     height, width, _ = img.shape
     total_pixels = width * height
 
@@ -101,6 +163,7 @@ def Decoder(Source):
 def MetaStego():
     token = "<YOUR_APYHUB_TOKEN_HERE>"  # Place your token here
     metadata = None  # To store the metadata
+    validator = InputValidator()
 
     while True:
         print("--Welcome to MetaStego--")
@@ -109,37 +172,60 @@ def MetaStego():
         print("3: Send Metadata")
         print("4: Exit")
 
-        func = input("Choose an option: ")
-
-        if func == '1':
-            print("Enter Source Image Path")
-            src = input()
-            print("Enter Message to Hide")
-            message = input()
-            print("Enter Destination Image Path")
-            dest = input()
-            print("Encoding...")
-            metadata = Encoder(src, message, dest)  # Store the returned metadata
-
-        elif func == '2':
-            print("Enter Source Image Path")
-            src = input()
-            print("Decoding...")
-            Decoder(src)
-
-        elif func == '3':
-            if metadata is None:
-                print("ERROR: No metadata available. Please encode an image first.")
-            else:
-                print("Sending metadata...")
-                send_metadata(metadata, token)
-
-        elif func == '4':
-            print("Exiting MetaStego...")
-            break
+        try :
+            func = input("Choose an option: ")
+
+            if func == '1':
+                # Get and validate source image
+                src = get_validated_input(
+                    "Enter Source Image Path: ",
+                    validator.validate_image_path,
+                    "Please provide a valid image file path"
+                )
+                
+                # Get and validate message
+                message = get_validated_input(
+                    "Enter Message to Hide: ",
+                    validator.validate_message,
+                    "Please provide a valid message"
+                )
+                
+                # Get and validate destination
+                dest = get_validated_input(
+                    "Enter Destination Image Path: ",
+                    validator.validate_destination_path,
+                    "Please provide a valid destination path"
+                )
+
+                print("Encoding...")
+                metadata = Encoder(src, message, dest)  # Store the returned metadata
+
+            elif func == '2':
+                src = get_validated_input(
+                    "Enter Source Image Path: ",
+                    validator.validate_image_path,
+                    "Please provide a valid image file path"
+                )
+
+                print("Decoding...")
+                Decoder(src)
+
+            elif func == '3':
+                if metadata is None:
+                    print("ERROR: No metadata available. Please encode an image first.")
+                else:
+                    print("Sending metadata...")
+                    send_metadata(metadata, token)
+
+            elif func == '4':
+                print("Exiting MetaStego...")
+                break
 
-        else:
-            print("ERROR: Invalid option chosen")
+            else:
+                print("ERROR: Invalid option. Please choose 1-4.")
+        except Exception as e:
+            print(f"An error occurred: {str(e)}")
+            print("Please try again.")
 
 if __name__ == "__main__":
     MetaStego()