From 499fd03e2e8c1e4b0ef67a6988db4ed7ff645ed8 Mon Sep 17 00:00:00 2001
From: Wassilios Lytras <chadgates@users.noreply.github.com>
Date: Wed, 1 May 2024 03:04:46 +0200
Subject: [PATCH] Extending Partner with Signature Algo and pass the setting to
 signing function.

* Fix github action build fail due to:
https://stackoverflow.com/questions/71673404/importerror-cannot-import-name-unicodefun-from-click

* Added partner setting to force canonicalize binary.

* Formatted with black

* https://github.com/abhishek-ram/pyas2-lib/issues/60
Extending Partner with Signature Algo and pass the setting to signing function.
---
 pyas2lib/as2.py                 | 21 +++++++++++++++++++--
 pyas2lib/constants.py           |  4 ++++
 pyas2lib/tests/test_advanced.py |  3 +++
 3 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/pyas2lib/as2.py b/pyas2lib/as2.py
index 12ebac7..338084d 100644
--- a/pyas2lib/as2.py
+++ b/pyas2lib/as2.py
@@ -28,6 +28,7 @@
     MDN_CONFIRM_TEXT,
     MDN_FAILED_TEXT,
     MDN_MODES,
+    SIGNATUR_ALGORITHMS,
     SYNCHRONOUS_MDN,
 )
 from pyas2lib.exceptions import (
@@ -179,6 +180,9 @@ class Partner:
 
     :param canonicalize_as_binary: force binary canonicalization for this partner
 
+    :param sign_alg: The signing algorithm to be used for generating the
+        signature. (default `rsassa_pkcs1v15`)
+
     """
 
     as2_name: str
@@ -197,6 +201,7 @@ class Partner:
     mdn_confirm_text: str = MDN_CONFIRM_TEXT
     ignore_self_signed: bool = True
     canonicalize_as_binary: bool = False
+    sign_alg: str = "rsassa_pkcs1v15"
 
     def __post_init__(self):
         """Run the post initialisation checks for this class."""
@@ -225,6 +230,12 @@ def __post_init__(self):
                 f"must be one of {DIGEST_ALGORITHMS}"
             )
 
+        if self.sign_alg and self.sign_alg not in SIGNATUR_ALGORITHMS:
+            raise ImproperlyConfigured(
+                f"Unsupported Signature Algorithm {self.sign_alg}, "
+                f"must be one of {SIGNATUR_ALGORITHMS}"
+            )
+
     def load_verify_cert(self):
         """Load the verification certificate of the partner and returned the parsed cert."""
         if self.validate_certs:
@@ -466,7 +477,10 @@ def build(
             )
             del signature["MIME-Version"]
             signature_data = sign_message(
-                mic_content, self.digest_alg, self.sender.sign_key
+                mic_content,
+                self.digest_alg,
+                self.sender.sign_key,
+                self.receiver.sign_alg,
             )
             signature.set_payload(signature_data)
             encoders.encode_base64(signature)
@@ -865,7 +879,10 @@ def build(
             del signature["MIME-Version"]
 
             signed_data = sign_message(
-                canonicalize(self.payload), self.digest_alg, message.receiver.sign_key
+                canonicalize(self.payload),
+                self.digest_alg,
+                message.receiver.sign_key,
+                message.sender.sign_alg,
             )
             signature.set_payload(signed_data)
             encoders.encode_base64(signature)
diff --git a/pyas2lib/constants.py b/pyas2lib/constants.py
index 53e6c1f..b5d4de2 100644
--- a/pyas2lib/constants.py
+++ b/pyas2lib/constants.py
@@ -28,3 +28,7 @@
     "aes_192_cbc",
     "aes_256_cbc",
 )
+SIGNATUR_ALGORITHMS = (
+    "rsassa_pkcs1v15",
+    "rsassa_pss",
+)
diff --git a/pyas2lib/tests/test_advanced.py b/pyas2lib/tests/test_advanced.py
index 0bc6db7..f010ecd 100644
--- a/pyas2lib/tests/test_advanced.py
+++ b/pyas2lib/tests/test_advanced.py
@@ -334,6 +334,9 @@ def test_partner_checks(self):
         with self.assertRaises(ImproperlyConfigured):
             as2.Partner("a partner", mdn_digest_alg="xyz")
 
+        with self.assertRaises(ImproperlyConfigured):
+            as2.Partner("a partner", sign_alg="xyz")
+
     def test_message_checks(self):
         """Test the checks and other features of Message."""
         msg = as2.Message()