Repository: MoshiBin/ssdpy

Issue:
No support for custom fields
SSDP allows custom fields to be send by the server. However ssdpy does not support them. 

See https://web.archive.org/web/20150905102426/http://upnp.org/specs/arch/UPnP-arch-DeviceArchitecture-v1.1.pdf chapter 1.1.3, SSDP header field extensions. 

Comments:

Patch:
diff --git a/CHANGES.md b/CHANGES.md
index 6617b09..3583296 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -4,6 +4,7 @@
 - Fixed an issue where `ssdpy.client.discover()` was using wrong syntax.
 - Changed the exception raised by `ssdpy.compat.if_nametoindex()` to be the same as in Python 3 (OSError).
 - Added tests for `ssdpy.client`, `ssdpy.compat` and created more tests for `ssdpy.server` to increase coverage.
+- Added support for custom fields in NOTIFY. Pass `extra_fields={"field": "value"}` to `ssdpy.SSDPServer` or pass `-e|--extra-field NAME VALUE` to `ssdpy-server`.
 
 ### 0.3.0
 (2020-08-10)
diff --git a/ssdpy/cli/server.py b/ssdpy/cli/server.py
index bbbb22d..b7f280f 100644
--- a/ssdpy/cli/server.py
+++ b/ssdpy/cli/server.py
@@ -11,17 +11,11 @@ logging.basicConfig()
 
 def parse_args(argv):
     parser = argparse.ArgumentParser(description="Start an SSDP server")
-    parser.add_argument(
-        "-V", "--version", action="version", version="%(prog)s {}".format(VERSION)
-    )
+    parser.add_argument("-V", "--version", action="version", version="%(prog)s {}".format(VERSION))
     parser.add_argument("-v", "--verbose", help="Be more verbose", action="store_true")
     proto_group = parser.add_mutually_exclusive_group()
-    proto_group.add_argument(
-        "-4", "--ipv4", help="Listen on IPv4 (default: True)", action="store_true"
-    )
-    proto_group.add_argument(
-        "-6", "--ipv6", help="Listen on IPv6 instead of IPv4", action="store_true"
-    )
+    proto_group.add_argument("-4", "--ipv4", help="Listen on IPv4 (default: True)", action="store_true")
+    proto_group.add_argument("-6", "--ipv6", help="Listen on IPv6 instead of IPv4", action="store_true")
     parser.add_argument("usn", help="Unique server name", nargs=1)
     parser.add_argument(
         "-t",
@@ -52,11 +46,22 @@ def parse_args(argv):
         "--address",
         help="Address of the interface to listen on. Only valid for IPv4.",
     )
+    parser.add_argument(
+        "-e",
+        "--extra-field",
+        action="append",
+        nargs=2,
+        metavar=("NAME", "VALUE"),
+        help="Extra fields to pass in NOTIFY packets. Pass multiple times for multiple extra headers",
+    )
     return parser.parse_args(argv)
 
 
 def main(argv=None):
     args = parse_args(argv)
+    extra_fields = None
+    if args.extra_field is not None:
+        extra_fields = dict(args.extra_field)
 
     if args.ipv6:
         proto = "ipv6"
@@ -76,6 +81,7 @@ def main(argv=None):
         max_age=args.max_age,
         al=args.location,
         location=args.location,
+        extra_fields=extra_fields,
     )
 
     logger = logging.getLogger("ssdpy.server")
diff --git a/ssdpy/protocol.py b/ssdpy/protocol.py
index a3b24e3..4961366 100644
--- a/ssdpy/protocol.py
+++ b/ssdpy/protocol.py
@@ -29,7 +29,7 @@ def create_msearch_payload(host, st, mx=1):
     return data.encode("utf-8")
 
 
-def create_notify_payload(host, nt, usn, location=None, al=None, max_age=None):
+def create_notify_payload(host, nt, usn, location=None, al=None, max_age=None, extra_fields=None):
     """
     Create a NOTIFY packet using the given parameters.
     Returns a bytes object containing a valid NOTIFY request.
@@ -63,6 +63,8 @@ def create_notify_payload(host, nt, usn, location=None, al=None, max_age=None):
     :param max_age: Amount of time in seconds that the NOTIFY packet should be cached by clients receiving it. In UPnP, this header is required.
     :type max_age: int
 
+    :param extra_fields: Extra header fields to send. UPnP SSDP section 1.1.3 allows for extra vendor-specific fields to be sent in the NOTIFY packet. According to the spec, the field names MUST be in the format of `token`.`domain-name`, for example `myheader.philips.com`. SSDPy, however, does not check this. Normally, headers should be in ASCII - but this function does not enforce that.
+
     :return: A bytes object containing the generated NOTIFY payload.
     """
     if max_age is not None and not isinstance(max_age, int):
@@ -80,5 +82,8 @@ def create_notify_payload(host, nt, usn, location=None, al=None, max_age=None):
         data += "AL:{}\r\n".format(al)
     if max_age is not None:
         data += "Cache-Control:max-age={}\r\n".format(max_age)
+    if extra_fields is not None:
+        for field, value in extra_fields.items():
+            data += "{}:{}\r\n".format(field, value)
     data += "\r\n"
     return data.encode("utf-8")
diff --git a/ssdpy/server.py b/ssdpy/server.py
index 1a2f026..28d789d 100644
--- a/ssdpy/server.py
+++ b/ssdpy/server.py
@@ -42,6 +42,8 @@ class SSDPServer(object):
     :type location: str
     :param al: Canonical URL of the service, but only supported in the IETF version of SSDP. Should be the same as ``location``.
     :type al: str
+    :param extra_fields: Extra header fields to send. UPnP SSDP section 1.1.3 allows for extra vendor-specific fields to be sent in the NOTIFY packet. According to the spec, the field names MUST be in the format of `token`.`domain-name`, for example `myheader.philips.com`. SSDPy, however, does not check this and allows any field name - as long as it's ASCII.
+    :type extra_fields: dict
     """
 
     def __init__(
@@ -55,6 +57,7 @@ class SSDPServer(object):
         max_age=None,
         location=None,
         al=None,
+        extra_fields=None,
     ):
         allowed_protos = ("ipv4", "ipv6")
         if proto not in allowed_protos:
@@ -67,6 +70,14 @@ class SSDPServer(object):
         self.max_age = max_age
         self._iface = iface
 
+        self._extra_fields = {}
+        if extra_fields is not None:
+            for field, value in extra_fields.items():
+                try:
+                    self._extra_fields[field.encode("ascii")] = value.encode("ascii")
+                except (UnicodeDecodeError, UnicodeEncodeError):
+                    raise ValueError("Invalid value for extra_field: %s=%s is not ASCII", field, value)
+
         if proto == "ipv4":
             self._af_type = socket.AF_INET
             self._broadcast_ip = ipv4_multicast_ip
@@ -125,7 +136,13 @@ class SSDPServer(object):
             logger.info("Received qualifying M-SEARCH from {}".format(address))
             logger.debug("M-SEARCH data: {}".format(headers))
             notify = create_notify_payload(
-                self._broadcast_ip, self.device_type, self.usn, self.location, self.al, self.max_age,
+                host=self._broadcast_ip,
+                nt=self.device_type,
+                usn=self.usn,
+                location=self.location,
+                al=self.al,
+                max_age=self.max_age,
+                extra_fields=self._extra_fields,
             )
             logger.debug("Created NOTIFY: {}".format(notify))
             try:


Test patch:
diff --git a/tests/test_cli_server.py b/tests/test_cli_server.py
index 77b4b5f..bd221d7 100644
--- a/tests/test_cli_server.py
+++ b/tests/test_cli_server.py
@@ -35,6 +35,7 @@ def test_ssdpserver_init(mocker):
         max_age=None,
         port=1900,
         proto="ipv4",
+        extra_fields=None,
     )
 
 
@@ -51,6 +52,7 @@ def test_ssdpserver_init_with_ipv6(mocker):
         max_age=None,
         port=1900,
         proto="ipv6",
+        extra_fields=None,
     )
 
     mocker.patch.object(server_cli, "SSDPServer")
@@ -65,6 +67,7 @@ def test_ssdpserver_init_with_ipv6(mocker):
         max_age=None,
         port=1900,
         proto="ipv6",
+        extra_fields=None,
     )
 
 
@@ -86,6 +89,9 @@ def test_ssdpserver_init_with_args(mocker):
             "test-device",
             "--max-age",
             "0",
+            "-e",
+            "test-field",
+            "foo"
         )
     )
     server_cli.SSDPServer.assert_called_once_with(
@@ -98,4 +104,5 @@ def test_ssdpserver_init_with_args(mocker):
         max_age=0,
         port=0,
         proto="ipv6",
+        extra_fields={"test-field": "foo"},
     )
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index f578525..67cc065 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -39,9 +39,7 @@ def test_notify_location():
 
 
 def test_notify_al():
-    data = create_notify_payload(
-        "239.255.255.250:1900", "testdevice", "ssdpy-test", al="http://localhost"
-    )
+    data = create_notify_payload("239.255.255.250:1900", "testdevice", "ssdpy-test", al="http://localhost")
     data_headers = parse_headers(data)
     assert data_headers.get("host") == "239.255.255.250:1900"
     assert data_headers.get("nt") == "testdevice"
@@ -51,9 +49,7 @@ def test_notify_al():
 
 
 def test_notify_age():
-    data = create_notify_payload(
-        "239.255.255.250:1900", "testdevice", "ssdpy-test", max_age=999
-    )
+    data = create_notify_payload("239.255.255.250:1900", "testdevice", "ssdpy-test", max_age=999)
     data_headers = parse_headers(data)
     assert data_headers.get("host") == "239.255.255.250:1900"
     assert data_headers.get("nt") == "testdevice"
@@ -65,3 +61,16 @@ def test_notify_age():
 def test_notify_edge_cases():
     with pytest.raises(ValueError):
         create_notify_payload("x", "y", "z", max_age="not-a-number")
+
+
+def test_notify_extra_fields():
+    data = create_notify_payload(
+        "239.255.255.250:1900",
+        "testdevice",
+        "ssdpy-test",
+        extra_fields={"test-header": "test-value", "test-header.domain.com": "test-value2"},
+    )
+    data_headers = parse_headers(data)
+    assert data_headers.get("test-header") == "test-value"
+    assert data_headers.get("test-header.domain.com") == "test-value2"
+    assert data_headers.get("non-existant-header") is None
diff --git a/tests/test_server.py b/tests/test_server.py
index edbeb5d..c16b507 100644
--- a/tests/test_server.py
+++ b/tests/test_server.py
@@ -67,3 +67,12 @@ def test_server_bind_address_and_iface_ipv6():
     except OSError as e:
         if e.errno != errno.ENOPROTOOPT:  # Protocol not supported
             raise
+
+
+def test_server_extra_fields():
+    SSDPServer("test-server", extra_fields={"test-field": "foo", "test-field2": "bar"})
+
+
+def test_server_extra_fields_non_ascii():
+    with pytest.raises(ValueError):
+        SSDPServer("test-server", extra_fields={"invalid-field™": "foo"})