Refactor examples for security and robustness

Co-authored-by: notapalindrome <notapalindrome@proton.me>

Author: cursoragent

python-examples/djvu-pdf-example.py

@@ -5,6 +5,8 @@
 import fnmatch
 import os
 import subprocess
+import shutil
+from pathlib import Path
 # global variables (change to suit your needs)
 inputfolderpath = '~'  # set to import folder path
 outputpath = '~'  # set to output folder (must exist)
@@ -26,32 +28,67 @@ def find_files(directory, pattern):
     for filename in find_files(inputfolderpath, '*.djvu'):
         print(f"[*] Processing DJVU to PDF for {filename}...")
         i = i + 1
-        inputfull = inputfolderpath+filename
-        outputfilename = filename[:-4]+i+'pdf'  # make filename unique
-        outputfilepath = outputpath
-        p = subprocess.Popen(["djvu2pdf", inputfull], stdout=subprocess.PIPE)
+        inputfull = os.path.join(inputfolderpath, filename)
+        # Validate that the file exists and is a regular file
+        if not os.path.isfile(inputfull):
+            print(f"[!] Skipping {filename} - not a valid file")
+            continue
+        outputfilename = f"{filename[:-5]}_{i}.pdf"  # make filename unique
+        outputfilepath = os.path.join(outputpath, outputfilename)
+        # Use list for subprocess to avoid shell injection
+        p = subprocess.Popen(
+            ["djvu2pdf", inputfull],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE
+        )
         output, err = p.communicate()
-        subprocess.call(["mv", outputfilename, outputfilepath])
+        # Use shutil.move instead of shell command for better security
+        if p.returncode == 0 and os.path.exists(outputfilename):
+            shutil.move(outputfilename, outputfilepath)
         print('[-] Processing finished for %s' % filename)
     print(f"[--] processed {i} file(s) [--]")
     exit('\n\"Sanity is madness put to good uses.\" - George Santayana\n')
 
 elif operationtype == '2':
     filename = input('What filename to process? (leave blank for example): ')
-    if 'djvu' in filename:
+    if filename and 'djvu' in filename:
+        # Validate filename to prevent path traversal
+        safe_path = Path(filename).resolve()
+        if not safe_path.is_file() or not str(safe_path).endswith('.djvu'):
+            print('[!] Invalid file or not a .djvu file')
+            exit('Invalid input')
         print('Processing DJVU to PDF...')
-        p = subprocess.Popen(["djvu2pdf", filename], stdout=subprocess.PIPE)
+        p = subprocess.Popen(
+            ["djvu2pdf", str(safe_path)],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE
+        )
         output, err = p.communicate()
-        print('Processing finished')
-        exit('Completed sucessfully')
+        if p.returncode == 0:
+            print('Processing finished')
+            exit('Completed successfully')
+        else:
+            print(f'[!] Error processing file: {err.decode() if err else "Unknown error"}')
+            exit('Failed')
     else:
         print('No djvu file to process, running sample')
         print('Processing DJVU to PDF...')
-        p = subprocess.Popen(["djvu2pdf", "assets/example.djvu"],
-                             stdout=subprocess.PIPE)
+        sample_file = Path("assets/example.djvu")
+        if not sample_file.is_file():
+            print('[!] Sample file not found')
+            exit('Sample file missing')
+        p = subprocess.Popen(
+            ["djvu2pdf", str(sample_file)],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE
+        )
         output, err = p.communicate()
-        print('Processing finished')
-        exit('Completed sucessfully')
+        if p.returncode == 0:
+            print('Processing finished')
+            exit('Completed successfully')
+        else:
+            print(f'[!] Error: {err.decode() if err else "Unknown error"}')
+            exit('Failed')
 
 
 elif operationtype == '':

python-examples/flask-example.py

@@ -39,26 +39,57 @@ def hello_world():
 @app.route("/upload", methods=["POST"])
 def upload_csv() -> str:
     """Upload CSV example."""
+    if "file" not in request.files:
+        return jsonify({"status": HTTPStatus.BAD_REQUEST, "message": "No file provided"}), 400
+    
     submitted_file = request.files["file"]
-    if submitted_file and allowed_filename(submitted_file.filename):
-        filename = secure_filename(submitted_file.filename)
-        directory = os.path.join(app.config["UPLOAD_FOLDER"])
-        if not os.path.exists(directory):
-            os.mkdir(directory)
-        basedir = os.path.abspath(os.path.dirname(__file__))
-        submitted_file.save(
-            os.path.join(basedir, app.config["UPLOAD_FOLDER"], filename)
-        )
-        out = {
-            "status": HTTPStatus.OK,
-            "filename": filename,
-            "message": f"{filename} saved successful.",
-        }
-        return jsonify(out)
+    
+    if not submitted_file or not submitted_file.filename:
+        return jsonify({"status": HTTPStatus.BAD_REQUEST, "message": "No file selected"}), 400
+    
+    if not allowed_filename(submitted_file.filename):
+        return jsonify({"status": HTTPStatus.BAD_REQUEST, "message": "File type not allowed"}), 400
+    
+    filename = secure_filename(submitted_file.filename)
+    
+    # Additional security check: ensure filename is not empty after sanitization
+    if not filename:
+        return jsonify({"status": HTTPStatus.BAD_REQUEST, "message": "Invalid filename"}), 400
+    
+    basedir = os.path.abspath(os.path.dirname(__file__))
+    upload_folder = os.path.abspath(os.path.join(basedir, app.config["UPLOAD_FOLDER"]))
+    
+    # Create directory with secure permissions if it doesn't exist
+    if not os.path.exists(upload_folder):
+        os.makedirs(upload_folder, mode=0o755, exist_ok=True)
+    
+    # Construct full path and verify it's within the upload directory (prevent path traversal)
+    file_path = os.path.abspath(os.path.join(upload_folder, filename))
+    
+    if not file_path.startswith(upload_folder):
+        return jsonify({"status": HTTPStatus.BAD_REQUEST, "message": "Invalid file path"}), 400
+    
+    # Limit file size (optional but recommended)
+    # submitted_file.seek(0, os.SEEK_END)
+    # file_size = submitted_file.tell()
+    # submitted_file.seek(0)
+    # if file_size > MAX_FILE_SIZE:
+    #     return jsonify({"status": HTTPStatus.BAD_REQUEST, "message": "File too large"}), 400
+    
+    submitted_file.save(file_path)
+    
+    out = {
+        "status": HTTPStatus.OK,
+        "filename": filename,
+        "message": f"{filename} saved successfully.",
+    }
+    return jsonify(out)
 
 
 if __name__ == "__main__":
     app.config["UPLOAD_FOLDER"] = "flaskme/"
-    app.run(port=6969, debug=True)
+    # Debug mode disabled for security - use environment variable to enable in development
+    debug_mode = os.environ.get("FLASK_DEBUG", "False").lower() == "true"
+    app.run(port=6969, debug=debug_mode)
 
 # curl -X POST localhost:6969/upload -F file=@"assets/archive_name.tar.gz" -i

python-examples/pickle_load-example.py

@@ -1,8 +1,28 @@
 # pickle load example
+# WARNING: pickle.load() can execute arbitrary code and should only be used
+# with trusted data. For untrusted data, use safer alternatives like JSON.
+# See: https://docs.python.org/3/library/pickle.html#module-pickle
 import pickle
 import random
+import os
 
-with open('assets/discordia.pkl', 'rb') as f:
+# Only load pickle files from trusted sources in trusted locations
+pickle_file = 'assets/discordia.pkl'
+
+# Verify the file exists and is in the expected location
+if not os.path.exists(pickle_file):
+    raise FileNotFoundError(f"Pickle file not found: {pickle_file}")
+
+# Resolve to absolute path to prevent path traversal
+pickle_file = os.path.abspath(pickle_file)
+expected_dir = os.path.abspath('assets')
+
+if not pickle_file.startswith(expected_dir):
+    raise ValueError("Pickle file must be in the assets directory")
+
+with open(pickle_file, 'rb') as f:
+    # SECURITY NOTE: This loads a pickle file that must be from a trusted source
+    # Never load pickle files from untrusted sources (user uploads, internet, etc.)
     discordia = pickle.load(f)