From f690555474a10920ce2dec93ccfb5fe292b486b6 Mon Sep 17 00:00:00 2001
From: Robin Krahl <robin@nitrokey.com>
Date: Fri, 3 May 2024 17:55:17 +0200
Subject: [PATCH] Add ignore_whitespace option to Base64.from_encoded

This patch makes it easy to ignore whitespace in Base64 input, e. g.
generated by the base64 tool without the --wrap=0 option.

Fixes: https://github.com/Nitrokey/nethsm-sdk-py/issues/108
---
 CHANGELOG.md       |  1 +
 Makefile           |  1 +
 nethsm/__init__.py | 40 +++++++++++++++++++++++++++++++++++-----
 3 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e2573d6..1f4e414 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,7 @@
 ### Features
 
 - Support key import from PEM files ([#99](https://github.com/Nitrokey/nethsm-sdk-py/issues/99))
+- Add `ignore_whitespace` option to `Base64.from_encoded` ([#108](https://github.com/Nitrokey/nethsm-sdk-py/issues/108))
 
 [All Changes](https://github.com/Nitrokey/nethsm-sdk-py/compare/v0.5.0...HEAD)
 
diff --git a/Makefile b/Makefile
index 203c3cf..a8eba58 100644
--- a/Makefile
+++ b/Makefile
@@ -66,4 +66,5 @@ nethsm-client: nethsm-api.yaml
 
 .PHONY: test
 test:
+	$(PYTHON3_VENV) -m doctest nethsm/__init__.py
 	$(PYTHON3_VENV) -m pytest --cov nethsm --cov-report=xml $(PYTEST_FLAGS)
diff --git a/nethsm/__init__.py b/nethsm/__init__.py
index d8a997a..50c99f5 100644
--- a/nethsm/__init__.py
+++ b/nethsm/__init__.py
@@ -14,6 +14,7 @@
 import contextlib
 import enum
 import json
+import string
 from base64 import b64decode, b64encode
 from dataclasses import dataclass
 from datetime import datetime, timezone
@@ -205,14 +206,43 @@ def decode(self) -> bytes:
         return b64decode(self.data)
 
     @classmethod
-    def from_encoded(cls, data: Union[bytes, str]) -> "Base64":
+    def from_encoded(
+        cls, data: Union[bytes, str], ignore_whitespace: bool = False
+    ) -> "Base64":
+        """
+        >>> Base64.from_encoded("dGVzdAo=")
+        Base64(data='dGVzdAo=')
+        >>> Base64.from_encoded(b"dGVzdAo=")
+        Base64(data='dGVzdAo=')
+        >>> Base64.from_encoded("dGV zdAo=")
+        Traceback (most recent call last):
+            ...
+        ValueError: Invalid base64 data: Non-base64 digit found: dGV zdAo=
+        >>> Base64.from_encoded(b"dGV zdAo=")
+        Traceback (most recent call last):
+            ...
+        ValueError: Invalid base64 data: Non-base64 digit found: dGV zdAo=
+        >>> Base64.from_encoded("dGV zdAo=", ignore_whitespace=True)
+        Base64(data='dGVzdAo=')
+        >>> Base64.from_encoded(b"dGV zdAo=", ignore_whitespace=True)
+        Base64(data='dGVzdAo=')
+        """
+        if ignore_whitespace:
+            if isinstance(data, bytes):
+                data = data.translate(None, delete=string.whitespace.encode())
+            else:
+                data = data.translate(str.maketrans("", "", string.whitespace))
+
         try:
             b64decode(data, validate=True)
+        except binascii.Error as e:
             if isinstance(data, bytes):
-                data = data.decode()
-            return cls(data=data)
-        except binascii.Error:
-            raise ValueError(f"Invalid base64 data: {data!r}")
+                data = data.decode(errors="replace")
+            raise ValueError(f"Invalid base64 data: {e}: {data}") from None
+
+        if isinstance(data, bytes):
+            data = data.decode()
+        return cls(data=data)
 
     @classmethod
     def encode(cls, data: bytes) -> "Base64":