UDP/IPv6 support added to asyncio-based transport

3 files changed, 71 insertions, 31 deletions

diff --git a/pysnmp/carrier/asyncio/dgram/base.py b/pysnmp/carrier/asyncio/dgram/base.py
index 25b64a2..a37aa26 100644
--- a/pysnmp/carrier/asyncio/dgram/base.py
+++ b/pysnmp/carrier/asyncio/dgram/base.py
@@ -39,6 +39,8 @@ loop = asyncio.get_event_loop()
 
 class DgramAsyncioProtocol(asyncio.DatagramProtocol, AbstractAsyncioTransport):
     """Base Asyncio datagram Transport, to be used with AsyncioDispatcher"""
+    sockFamily = None
+    addressType = lambda x: x
     transport = None
 
     def __init__(self, *args, **kwargs):
@@ -58,13 +60,41 @@ class DgramAsyncioProtocol(asyncio.DatagramProtocol, AbstractAsyncioTransport):
             debug.logger & debug.flagIO and debug.logger('connection_made: transportAddress %r outgoingMessage %s' %
                                                          (transportAddress, debug.hexdump(outgoingMessage)))
             try:
-                self.transport.sendto(outgoingMessage, transportAddress)
+                self.transport.sendto(outgoingMessage, self.normalizeAddress(transportAddress))
             except Exception:
                 raise error.CarrierError(';'.join(traceback.format_exception(*sys.exc_info())))
 
     def connection_lost(self, exc):
         debug.logger & debug.flagIO and debug.logger('connection_lost: invoked')
 
+    # AbstractAsyncioTransport API
+
+    def openClientMode(self, iface=None):
+        try:
+            c = loop.create_datagram_endpoint(
+                lambda: self, local_addr=iface, family=self.sockFamily
+            )
+            self._lport = asyncio.async(c)
+        except Exception:
+            raise error.CarrierError(';'.join(traceback.format_exception(*sys.exc_info())))
+        return self
+
+    def openServerMode(self, iface):
+        try:
+            c = loop.create_datagram_endpoint(
+                lambda: self, local_addr=iface, family=self.sockFamily
+            )
+            self._lport = asyncio.async(c)
+        except Exception:
+            raise error.CarrierError(';'.join(traceback.format_exception(*sys.exc_info())))
+        return self
+
+    def closeTransport(self):
+        self._lport.cancel()
+        if self.transport is not None:
+            self.transport.close()
+        AbstractAsyncioTransport.closeTransport(self)
+
     def sendMessage(self, outgoingMessage, transportAddress):
         debug.logger & debug.flagIO and debug.logger('sendMessage: %s transportAddress %r outgoingMessage %s' % (
             (self.transport is None and "queuing" or "sending"),
@@ -74,12 +104,11 @@ class DgramAsyncioProtocol(asyncio.DatagramProtocol, AbstractAsyncioTransport):
             self._writeQ.append((outgoingMessage, transportAddress))
         else:
             try:
-                self.transport.sendto(outgoingMessage, transportAddress)
+                self.transport.sendto(outgoingMessage, self.normalizeAddress(transportAddress))
             except Exception:
                 raise error.CarrierError(';'.join(traceback.format_exception(*sys.exc_info())))
 
-    def closeTransport(self):
-        if self.transport is not None:
-            self.transport.close()
-        AbstractAsyncioTransport.closeTransport(self)
-
+    def normalizeAddress(self, transportAddress):
+        if not isinstance(transportAddress, self.addressType):
+            transportAddress = self.addressType(transportAddress)
+        return transportAddress
diff --git a/pysnmp/carrier/asyncio/dgram/udp.py b/pysnmp/carrier/asyncio/dgram/udp.py
index 328423b..653c83c 100644
--- a/pysnmp/carrier/asyncio/dgram/udp.py
+++ b/pysnmp/carrier/asyncio/dgram/udp.py
@@ -25,8 +25,7 @@
 # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF 
 # THE POSSIBILITY OF SUCH DAMAGE.
 #
-import sys
-import traceback
+import socket
 from pysnmp.carrier.base import AbstractTransportAddress
 from pysnmp.carrier.asyncio.dgram.base import DgramAsyncioProtocol
 from pysnmp.carrier import error
@@ -42,28 +41,7 @@ domainName = snmpUDPDomain = (1, 3, 6, 1, 6, 1, 1)
 class UdpTransportAddress(tuple, AbstractTransportAddress): pass
 
 class UdpAsyncioTransport(DgramAsyncioProtocol):
+    sockFamily = socket.AF_INET
     addressType = UdpTransportAddress
 
-    # AbstractAsyncioTransport API
-
-    def openClientMode(self, iface=('0.0.0.0', 0)):
-        try:
-            c = loop.create_datagram_endpoint(lambda: self, local_addr=iface)
-            self._lport = asyncio.async(c)
-        except Exception:
-            raise error.CarrierError(';'.join(traceback.format_exception(*sys.exc_info())))
-        return self
-
-    def openServerMode(self, iface=('0.0.0.0', 161)):
-        try:
-            c = loop.create_datagram_endpoint(lambda: self, local_addr=iface)
-            self._lport = asyncio.async(c)
-        except Exception:
-            raise error.CarrierError(';'.join(traceback.format_exception(*sys.exc_info())))
-        return self
-
-    def closeTransport(self):
-        self._lport.cancel()
-        DgramAsyncioProtocol.closeTransport(self)
-
 UdpTransport = UdpAsyncioTransport
diff --git a/pysnmp/carrier/asyncio/dgram/udp6.py b/pysnmp/carrier/asyncio/dgram/udp6.py
new file mode 100644
index 0000000..8d79174
--- /dev/null
+++ b/pysnmp/carrier/asyncio/dgram/udp6.py
@@ -0,0 +1,33 @@
+import socket
+from pysnmp.carrier.base import AbstractTransportAddress
+from pysnmp.carrier.asyncio.dgram.base import DgramAsyncioProtocol
+from pysnmp.carrier import error
+try:
+    import asyncio
+except ImportError:
+    import trollius as asyncio
+
+loop = asyncio.get_event_loop()
+
+domainName = snmpUDP6Domain = (1, 3, 6, 1, 2, 1, 100, 1, 2)
+
+class Udp6TransportAddress(tuple, AbstractTransportAddress): pass
+
+class Udp6AsyncioTransport(DgramAsyncioProtocol):
+    sockFamily = socket.has_ipv6 and socket.AF_INET6 or None
+    addressType = Udp6TransportAddress
+
+    def normalizeAddress(self, transportAddress):
+        if '%' in transportAddress[0]:  # strip zone ID
+            return self.addressType(
+                    (transportAddress[0].split('%')[0],
+                     transportAddress[1],
+                     0, # flowinfo
+                     0) # scopeid
+                )
+        else:
+            return self.addressType(
+                    (transportAddress[0], transportAddress[1], 0, 0)
+                )
+
+Udp6Transport = Udp6AsyncioTransport