Fixed a few bugs and broken feature implentation.

2025-12-28 13:12:40 -05:00
parent 653de8cbaa
commit 541e6424ea
6 changed files with 176 additions and 29 deletions
--- a/frontends/web/app.py
+++ b/frontends/web/app.py
@@ -254,10 +254,41 @@ def generate_qr_download(token):
        return f"Error generating QR code: {e}", 500


+#@app.route('/generate/download-key', methods=['POST'])
+#def download_key():
+#    """Download RSA key as password-protected PEM file."""
+#    key_pem = request.form.get('key_pem', '')
+
@app.route('/generate/download-key', methods=['POST'])
 def download_key():
    """Download RSA key as password-protected PEM file."""
    key_pem = request.form.get('key_pem', '')
+    password = request.form.get('key_password', '')
+
+    if not key_pem:
+        flash('No key to download', 'error')
+        return redirect(url_for('generate'))
+
+    if not password or len(password) < 8:
+        flash('Password must be at least 8 characters', 'error')
+        return redirect(url_for('generate'))
+
+    try:
+        private_key = load_rsa_key(key_pem.encode('utf-8'))
+        encrypted_pem = export_rsa_key_pem(private_key, password=password)
+
+        key_id = secrets.token_hex(4)
+        filename = f'stegasoo_key_{private_key.key_size}_{key_id}.pem'
+
+        return send_file(
+            io.BytesIO(encrypted_pem),
+            mimetype='application/x-pem-file',
+            as_attachment=True,
+            download_name=filename
+        )
+    except Exception as e:
+        flash(f'Error creating key file: {e}', 'error')
+        return redirect(url_for('generate'))


@app.route('/extract-key-from-qr', methods=['POST'])
--- a/frontends/web/templates/decode.html
+++ b/frontends/web/templates/decode.html
@@ -151,7 +151,7 @@
                                    <input type="file" name="rsa_key" class="form-control form-control-sm" id="rsaKeyInput" accept=".pem,.key,application/x-pem-file">
                                </div>
                                <div class="tab-pane fade" id="rsaQrTabDec" role="tabpanel">
-                                    <input type="file" name="rsa_key_qr" class="form-control form-control-sm" id="rsaKeyQrInput" accept="image/png,image/jpeg,image/gif,image/webp,.png,.jpg,.jpeg,.gif,.webp">
+                                  <input type="file" name="rsa_key_qr" class="form-control form-control-sm" id="rsaKeyQrInput" accept="image/*">
                                    <div class="form-text small">PNG, JPG, or other image of QR code</div>
                                </div>
                            </div>
--- a/frontends/web/templates/encode.html
+++ b/frontends/web/templates/encode.html
@@ -160,7 +160,7 @@
                                    <input type="file" name="rsa_key" class="form-control form-control-sm" id="rsaKeyInput" accept=".pem,.key,application/x-pem-file">
                                </div>
                                <div class="tab-pane fade" id="rsaQrTab" role="tabpanel">
-                                    <input type="file" name="rsa_key_qr" class="form-control form-control-sm" id="rsaKeyQrInput" accept="image/png,image/jpeg,image/gif,image/webp,.png,.jpg,.jpeg,.gif,.webp">
+                                    <input type="file" name="rsa_key_qr" class="form-control form-control-sm" id="rsaKeyQrInput" accept="image/*">
                                    <div class="form-text small">PNG, JPG, or other image of QR code</div>
                                </div>
                            </div>
@@ -189,7 +189,7 @@
                        <i class="bi bi-lock me-2"></i>Encode
                    </button>
                </form>
-                
+               
                <hr class="my-4">
                
                <div class="row text-center text-muted small">
--- a/requirements.txt
+++ b/requirements.txt
@@ -14,6 +14,7 @@ pyzbar>=0.1.9

 # Web Frontend (Flask)
 Flask>=3.0.0
+pyzbar>=0.1.8

 # API Frontend (FastAPI)
 fastapi>=0.109.0
--- a/src/stegasoo/qr_utils.py
+++ b/src/stegasoo/qr_utils.py
@@ -3,6 +3,11 @@ Stegasoo QR Code Utilities

 Functions for generating and reading QR codes containing RSA keys.
 Supports automatic compression for large keys.
+
+IMPROVEMENTS IN THIS VERSION:
+- Much more robust PEM normalization
+- Better handling of QR code extraction edge cases
+- Improved error messages
 """

 import io
@@ -72,47 +77,88 @@ def decompress_data(data: str) -> str:

 def normalize_pem(pem_data: str) -> str:
    """
-    Normalize PEM data to ensure proper formatting.
+    Normalize PEM data to ensure proper formatting for cryptography library.
    
-    Fixes common issues:
-    - Inconsistent line endings
+    The cryptography library is very particular about PEM formatting.
+    This function handles all common issues from QR code extraction:
+    - Inconsistent line endings (CRLF, LF, CR)
    - Missing newlines after header/before footer
-    - Extra whitespace
+    - Extra whitespace, tabs, multiple spaces
+    - Non-ASCII characters
+    - Incorrect base64 padding
+    - Malformed headers/footers
    
    Args:
-        pem_data: Raw PEM string
+        pem_data: Raw PEM string from QR code
        
    Returns:
-        Properly formatted PEM string
+        Properly formatted PEM string that cryptography library will accept
    """
    import re
    
-    # Normalize line endings
+    # Step 1: Normalize ALL line endings to \n
    pem_data = pem_data.replace('\r\n', '\n').replace('\r', '\n')
    
-    # Remove any leading/trailing whitespace
+    # Step 2: Remove leading/trailing whitespace
    pem_data = pem_data.strip()
    
-    # Extract header, content, and footer using regex
-    # Match patterns like -----BEGIN RSA PRIVATE KEY----- or -----BEGIN PRIVATE KEY-----
-    pattern = r'(-----BEGIN [^-]+-----)(.*?)(-----END [^-]+-----)'
-    match = re.search(pattern, pem_data, re.DOTALL)
+    # Step 3: Remove any non-ASCII characters (QR artifacts)
+    pem_data = ''.join(char for char in pem_data if ord(char) < 128)
+    
+    # Step 4: Extract header, content, and footer with flexible regex
+    # This handles variations like:
+    # - "PRIVATE KEY" vs "RSA PRIVATE KEY"
+    # - Extra spaces in headers
+    # - Missing spaces
+    pattern = r'(-----BEGIN[^-]*-----)(.*?)(-----END[^-]*-----)'
+    match = re.search(pattern, pem_data, re.DOTALL | re.IGNORECASE)
    
    if not match:
-        return pem_data  # Return as-is if not recognized
+        # Fallback: try even more permissive pattern
+        pattern = r'(-+BEGIN[^-]+-+)(.*?)(-+END[^-]+-+)'
+        match = re.search(pattern, pem_data, re.DOTALL | re.IGNORECASE)
+        
+        if not match:
+            # Last resort: return original if can't parse
+            return pem_data
    
-    header = match.group(1)
-    content = match.group(2)
-    footer = match.group(3)
+    header_raw = match.group(1).strip()
+    content_raw = match.group(2)
+    footer_raw = match.group(3).strip()
    
-    # Clean up the base64 content
-    # Remove all whitespace and rejoin with proper 64-char lines
-    content_clean = ''.join(content.split())
+    # Step 5: Normalize header and footer
+    # Standardize spacing and ensure proper format
+    header = re.sub(r'\s+', ' ', header_raw)
+    footer = re.sub(r'\s+', ' ', footer_raw)
    
-    # Split into 64-character lines (PEM standard)
+    # Ensure exactly 5 dashes on each side
+    header = re.sub(r'^-+', '-----', header)
+    header = re.sub(r'-+$', '-----', header)
+    footer = re.sub(r'^-+', '-----', footer)
+    footer = re.sub(r'-+$', '-----', footer)
+    
+    # Step 6: Clean the base64 content THOROUGHLY
+    # Remove ALL whitespace: spaces, tabs, newlines
+    # Keep only valid base64 characters: A-Z, a-z, 0-9, +, /, =
+    content_clean = ''.join(
+        char for char in content_raw 
+        if char.isalnum() or char in '+/='
+    )
+    
+    # Double-check: remove any remaining invalid characters
+    content_clean = re.sub(r'[^A-Za-z0-9+/=]', '', content_clean)
+    
+    # Step 7: Fix base64 padding
+    # Base64 strings must be divisible by 4
+    remainder = len(content_clean) % 4
+    if remainder:
+        content_clean += '=' * (4 - remainder)
+    
+    # Step 8: Split into 64-character lines (PEM standard)
    lines = [content_clean[i:i+64] for i in range(0, len(content_clean), 64)]
    
-    # Reconstruct PEM
+    # Step 9: Reconstruct with EXACT PEM formatting
+    # Format: header\ncontent_line1\ncontent_line2\n...\nfooter\n
    return header + '\n' + '\n'.join(lines) + '\n' + footer + '\n'


@@ -278,30 +324,45 @@ def extract_key_from_qr(image_data: bytes) -> Optional[str]:
    """
    Extract RSA key from QR code image, auto-decompressing if needed.
    
+    This function is more robust than the original, with better error handling
+    and PEM normalization.
+    
    Args:
        image_data: Image bytes containing QR code
        
    Returns:
        PEM-encoded RSA key string, or None if not found/invalid
    """
+    # Step 1: Read QR code
    qr_data = read_qr_code(image_data)
    
    if not qr_data:
        return None
    
-    # Auto-decompress if needed
+    # Step 2: Auto-decompress if needed
    try:
        if is_compressed(qr_data):
            key_pem = decompress_data(qr_data)
        else:
            key_pem = qr_data
-    except Exception:
+    except Exception as e:
+        # If decompression fails, try using data as-is
        key_pem = qr_data
    
-    # Validate it looks like a PEM key
-    if '-----BEGIN' in key_pem and '-----END' in key_pem:
-        # Normalize PEM format to fix potential line ending issues
+    # Step 3: Validate it looks like a PEM key
+    if '-----BEGIN' not in key_pem or '-----END' not in key_pem:
+        return None
+    
+    # Step 4: Aggressively normalize PEM format
+    # This is crucial - QR codes can introduce subtle formatting issues
+    try:
        key_pem = normalize_pem(key_pem)
+    except Exception as e:
+        # If normalization fails, return None rather than broken PEM
+        return None
+    
+    # Step 5: Final validation - ensure it still looks like PEM
+    if '-----BEGIN' in key_pem and '-----END' in key_pem:
        return key_pem
    
    return None
--- a/test_data/stegasoo_key_4096_5e663335.pem
+++ b/test_data/stegasoo_key_4096_5e663335.pem
@@ -0,0 +1,54 @@
+-----BEGIN ENCRYPTED PRIVATE KEY-----
+MIIJtTBfBgkqhkiG9w0BBQ0wUjAxBgkqhkiG9w0BBQwwJAQQJdD0f2FnF8tXObq2
+HeQj8QICCAAwDAYIKoZIhvcNAgkFADAdBglghkgBZQMEASoEEFKD/tvo6am/xKKS
+fiNtbagEgglQBJdTsd1JIjihIK+tcV+SbNJggJ0i7R0sh82GxZ21Oca2Ij4FndPU
+rwjhyv8977dibIwt1F6oJOkWgt/DLCFVMinQvJaKdKY2Jowgj42MfiRQlFnzXJhY
+GI1LHPg4/PWBNUIWKrOYOlVB+Nq4SffjQFlpmQGSxCjLwCNLZCG0ckxWBFrHg1g1
+R1LPnQikBEJ1xvtyMHELlyQia2JPDwvn29vhGtT5Jr9y4762R86RgqbelbB7H5wn
+4WG4b9agZERx9vwnF7NQEFpOOhe6CMjEsWdfSswAsUoz/zaHmVz2alCOlQYj1yJj
+vDPbHR9NZc1UtuH7g0pbEijUIto/PZcYhXPEvb1knwOA/JY7DuCmvW1t1rNsTSqk
+2L8kmjDlr2FDDcNvD2XLHVZzqp3F5jYLtXfkWpOH7rqkrvdqHeu+ve5jxCnesZ0D
+rDpcmpbEwqWx/W3slpZEqAdTrSgLcXXDi6OjIzAYDEzCHO/u6djDDKzYF1ziZNxq
+bq4ZogP4SfzaGehArnCbPIBIObQp8t2BuXk6veDmEHk4aPSSBbbjKhWXVSbposz4
+ZvespTu2Z4aIT+xb7Rj32fAjiy+IPEI7Mt/KtsV+W2F5CM+QQxWTOdUkt+3OuAJe
+VlgnZk4a7yHYLXbyqc/wpHPdD4EEKyCCBuT2lPwu+L/3XNNy8dWL/1y74PbUOyAW
+r5wfIalJZ43Zabvgl+LXxCUXrVRFMG1hASXupCY88uU1evvdBjd+anWTd/IpNHBC
+g6pvwnHQDeuf9KhzKIRvb2HqMeYM80yir6PMBcayZj6icKSZa9i2KKs6W4IVhS1p
+ZDZBbuP01GlwU3pAX+bX7HIBt9wPYYoabUjYDahsvLCKToK8rhLbHd//3qKOuIh2
+7T+DtouVTFu7ipuxaq+VqSAExU5gNXi9xh9fSbJwAf7E//LA9s6UBMTRfJOmC1Wv
+gyapSNqeATkvwFNmucTIXbaFTTlR+6WisgEO7eqT7F99k+tDoj/m8HoX84mcesqz
+t3zeR0A6L0bq0GAICxdkNMRMXZWuan34T7IvxjdtIsaUm3ReIDf68oW51107Wlts
+ZX2IE4P+vrAq4gR0Ra4L2NaDWDawZMIyEFAMRHxNE96TqZzvaNVZW3dOfn0YjRJH
+fuvRThuoDGKKM5NzVDuWQJM+PP3dR2I+wamiL4QEeP+czP5FQXxR2C5iwY03Ntcj
+ByAp1ZiLoGePEu3PGFIAocntyIy+UTKVMLfvqn1tX3VW19uF4J8eQnp0W4oqOAcZ
+DTV6gamXNHrJzI5qtlB5yBf1YZb1bxniLKCiihOyx1O3fY/y178gIePMXX1ZVpQZ
+PWdYlyDlw07tk5WnQxxAj4E6iNodlkhm9lfBFf+8GPgFe3esgPyID79KbS5UqN6D
+gpnJcV57vsbU8KkjZ1hYEHUCuyR3AWIQOGAjP0Ai/nJADtEF54UZbP6fnOPT6yJR
+olek4GiaEFV9SiSReIwKeTHiCZvpN1rMDnGLTn9p2bphOBM8mjBhKfE8Wy1LOYr9
+5HjJleAgtppgDh0dnKPc6kV0e+yHeQXDp0o1RC2J0awW4Oeqr65dJOoynARQ95n6
+UVlahI07BKqWZNRKcmJVvrWaQisDDLfWrvCaGYocTfOBEb9mpJzLZ3NrtE6UBxSj
+/caJH0y0dRBaDLJvH44RXK9hXVW0iRp09lpABID9AvUyAFc/G+aKTbxbHkhc6AwB
+pITCXPC+EMQ7Z4TcoRykU6+6EMsYNjvZ0l5xpsh5Pe7zsNeBtmBa8z//71ZkjsFi
+Ioy3dmD0ruWgkq0dlU2L0BfNr55tsCZUzfd9/u2/hE6Ye4edtsKKQJD6aqoMi8Nk
+qDI4t2GS1RHiCZ8hr4Ux5NXvKCFxD5913n5OY70BtMXKg/H/TwoTBqwzSH6fv1JZ
+mWUSdtS9hN3fcezkqDwfR8Dzgz6Aq8ewa2HBoqcZ9T551hEGwvyN9QnT0DzkaZNK
+VNwvTAHQ5Xs3lbS0X+Giu75nvHJMpKL70Z/aNX5IwobmfAi89jXaUMuGetcVbO98
+SL96j5AxFO6K0PczCgE8CHXJY62Sh/eGYF+Uc7DbRZROxgM035MYBQqa5U17W0/G
+h2Mf+qvfrH0jsvTwod9BRbYusnxp0E04+1Y7SdcQfbcbpafc2MAjnQGxU51KQiWf
+yZ4D6COBoT1j7eGc/fg6uFKClEH34I97vod89CMj0uJblCieYj+5+pz0aGCgL3yP
+6WZb5ogZQkq23p7lMZptmjW+OZGNt5bNEqNTAIhRB5jN1PnvJs81vzQo1rNmoJG6
+rokC3A6Mqic7MssU0B9nUXUA92LEB/YhimO/sccRshbBD2/TuY+KhQdApbU7NtJ3
+giyj/5JEwUmj4ecGXfxhxWYfPrnLG87hO1mogfp1ndCC0efbLR6u8Qb5vlz56luQ
+hSvE45gWcVjxo6hJasZHpoqq4aD4CqVLgCi0zSEgXhPS+vgo2CYpW3u5N5Kw3nJG
+WmcQOfGUXIoCLsFoiSoLNt5H5uPXi4+rcgi65pio2QwXpYfxlCZpHbEgyvzr+U85
+fiBNPwSvYnQx3DYqx/2mkIZPJO1pSGfDKy68OAnvOMUhQ7jASgmMjK0HeRSpT1E7
+n3+cUk1zJgDbu68laxj0xzU+iyJZr4hk05mmqVfux60WSv7NqurLgLQ3++CZ5XPu
+SSuYY89gBlbbl9GLlF8EcmsbqXfqYa1F+6A2bqFBe96jbVo7WEdNXJDuZZxwU2GU
+FgDo9tyLxnkGfv3XfSBmZDydltOQm2sgGIZ0EXczbso0F4BDeamolCgL6jhgVs0B
+rhJ2kooSEA8/MJMhzUVgRjqNUV6iCW+iFRtX5nD4rW/vODYpFKs/zlSQo5qq8P3/
+eKw7VFlcc/i2V7ZxA48WIvM9HsNsKs3sHCxEHUZHmT/8KTcHuY1LlUA8aE3UMyAB
+iqrpMQwn3x6G5UqLa/3IoxGYH9dYvoDjESVKm9CTZjbQdiCpYENsNiZ+TkjBBUwU
+m50oRjC8YWhqAHdJxDcbAiiH0zyDYrMgvozLbDpUMjye8wOV94ga4Pb681Qld1vW
+rFfytkJPYFCIP0uVrlEuAnfrcvymLAB/tMEbMeiEoFuoRfy2ra7taOeH6tpQcb1N
+18QSzGTAcerjkvrpJLxG/aGyzKQDFvnpbObvsH3XJQScTgjhoY3yXPI=
+-----END ENCRYPTED PRIVATE KEY-----