Small update

Signed-off-by: Lev Rusanov <30170278+JDM170@users.noreply.github.com>

main.py

@@ -4,19 +4,14 @@
 from sys import exit
 from re import search
 from hashlib import md5
-from os.path import exists, getsize
+from os.path import isfile, getsize
 from progress.bar import IncrementalBar
 
 hash_chunk_size = 65536
 hash_pattern = r"[a-zA-Z0-9]{32,}" # Паттерн поиска хэш-суммы
 
 
-def check_hash_file(fn):
-    if not exists(fn):
-        if search(hash_pattern, fn).group() is not None:
-            return fn
-        print(" ! Hash sum not found")
-        return False
+def read_hash_from_file(fn):
     with open(fn, "r") as f:
         fdata = f.read()
         found_hash = search(hash_pattern, fdata).group()
@@ -28,9 +23,6 @@ def check_hash_file(fn):
 
 
 def generate_md5(fn):
-    if not exists(fn):
-        print(" ! File not found")
-        return False
     hash_md5 = md5()
     with open(fn, "rb") as f:
         bar = IncrementalBar(" Generating hash sum...", max=int(getsize(fn) / hash_chunk_size), suffix='%(percent)d%%')
@@ -41,17 +33,28 @@ def generate_md5(fn):
     return hash_md5.hexdigest()
 
 
+def ask_input(message):
+    result = input(message)
+    if not isfile(result):
+        ask_input(message)
+    return result
+
+
 def main():
-    file1 = str(input("\n Input first file: "))
-    file2 = str(input(" Input hash sum (or file): "))
-    selected_mode = bool(int(input(" Select compare mode (0 or 1)\n 0 - hash sum; 1 - second file path: ")))
-    # print(" selected_mode:", selected_mode)
-    if (file1 != "") and (file2 != ""):
-        if (selected_mode is False and exists(file2) and not file2.endswith(".txt")) or (selected_mode is True and not exists(file2)):
-        # if (not exists(file2) and selected_mode is True):
-            print(" ! Selected invalid mode")
-            main()
-        print("\n Equals?:", generate_md5(file1) == (check_hash_file(file2) if selected_mode is False else generate_md5(file2)), "\n")
+    src_file = ask_input("\n Input file: ")
+    hash_input = input(" Input hash sum (or file): ")
+    if (hash_input != ""):
+        src_file_result = generate_md5(src_file)
+        hash_input_result = False
+        if isfile(hash_input):
+            if hash_input.endswith(".md5"):
+                hash_input_result = read_hash_from_file(hash_input)
+            else:
+                hash_input_result = generate_md5(hash_input)
+        else:
+            if search(hash_pattern, hash_input).group() is not None:
+                hash_input_result = hash_input
+        print("\n Equals?:", (src_file_result == hash_input_result), "\n")
     main()
 
 
@@ -59,5 +62,4 @@ if __name__ == '__main__':
     try:
         main()
     except KeyboardInterrupt:
-        if input("\n Stop comparing? (y/n): ").lower() == "y":
-            exit()
+        exit()