Direct conversion between byte string and long, using scary

ctypes.pythonapi interface.

svn path=/branches/tk274/; revision=4629

1 files changed, 54 insertions, 61 deletions

diff --git a/rpkid/rpki/ipaddrs.py b/rpkid/rpki/ipaddrs.py
index a192f92b..20317118 100644
--- a/rpkid/rpki/ipaddrs.py
+++ b/rpkid/rpki/ipaddrs.py
@@ -42,86 +42,69 @@ OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
 PERFORMANCE OF THIS SOFTWARE.
 """
 
-import socket, struct
-
-class v4addr(long):
+import socket
+import ctypes
+
+# Scary hack to let us use methods from the Python/C API
+
+PyLong_AsByteArray = ctypes.pythonapi._PyLong_AsByteArray
+PyLong_AsByteArray.argtypes = [ctypes.py_object,
+                               ctypes.c_char_p,
+                               ctypes.c_size_t,
+                               ctypes.c_int,
+                               ctypes.c_int]
+
+PyLong_FromByteArray = ctypes.pythonapi._PyLong_FromByteArray
+PyLong_FromByteArray.restype = ctypes.py_object
+PyLong_FromByteArray.argtypes = [ctypes.c_char_p,
+                                 ctypes.c_size_t,
+                                 ctypes.c_int,
+                                 ctypes.c_int]
+
+class addr(long):
   """
-  IPv4 address.
-
-  Derived from long, but supports IPv4 print syntax.
+  IP address.  This is a virtual class.
+  Derived from long, but supports IP print syntax.
+  Derived classes must define .bits and .af values
+  and may override .normalize_string().
   """
 
-  bits = 32
-  
   def __new__(cls, x):
-    """
-    Construct a v4addr object.
-    """
     if isinstance(x, unicode):
       x = x.encode("ascii")
     if isinstance(x, str):
-      return cls.from_bytes(socket.inet_pton(socket.AF_INET, ".".join(str(int(i)) for i in x.split("."))))
+      return cls.from_bytes(socket.inet_pton(cls.af, cls.normalize_string(x)))
     else:
       return long.__new__(cls, x)
 
+  @staticmethod
+  def normalize_string(s):
+    return s
+
   def to_bytes(self):
-    """
-    Convert a v4addr object to a raw byte string.
-    """
-    return struct.pack("!I", long(self))
+    b = ctypes.create_string_buffer(self.bits / 8)
+    PyLong_AsByteArray(self, b, len(b), 0, 1)
+    return b.raw
 
   @classmethod
   def from_bytes(cls, x):
-    """
-    Convert from a raw byte string to a v4addr object.
-    """
-    return cls(struct.unpack("!I", x)[0])
+    return cls(PyLong_FromByteArray(x, len(x), 0, 1))
 
   def __str__(self):
-    """
-    Convert a v4addr object to string format.
-    """
-    return socket.inet_ntop(socket.AF_INET, self.to_bytes())
-
-class v6addr(long):
-  """
-  IPv6 address.
+    b = self.to_bytes()
+    return socket.inet_ntop(self.af, b)
 
-  Derived from long, but supports IPv6 print syntax.
-  """
+class v4addr(addr):
+  bits = 32
+  af = socket.AF_INET
+  
+  @staticmethod
+  def normalize_string(s):
+    return ".".join(str(int(i)) for i in s.split("."))
 
+class v6addr(addr):
   bits = 128
-
-  def __new__(cls, x):
-    """
-    Construct a v6addr object.
-    """
-    if isinstance(x, unicode):
-      x = x.encode("ascii")
-    if isinstance(x, str):
-      return cls.from_bytes(socket.inet_pton(socket.AF_INET6, x))
-    else:
-      return long.__new__(cls, x)
-
-  def to_bytes(self):
-    """
-    Convert a v6addr object to a raw byte string.
-    """
-    return struct.pack("!QQ", long(self) >> 64, long(self) & 0xFFFFFFFFFFFFFFFF)
-
-  @classmethod
-  def from_bytes(cls, x):
-    """
-    Convert from a raw byte string to a v6addr object.
-    """
-    x = struct.unpack("!QQ", x)
-    return cls((x[0] << 64) | x[1])
-
-  def __str__(self):
-    """
-    Convert a v6addr object to string format.
-    """
-    return socket.inet_ntop(socket.AF_INET6, self.to_bytes())
+  af = socket.AF_INET6
 
 def parse(s):
   """
@@ -130,3 +113,13 @@ def parse(s):
   if isinstance(s, unicode):
     s = s.encode("ascii")
   return v6addr(s) if ":" in s else v4addr(s)
+
+if __name__ == "__main__":
+  def test(x):
+    y = parse(x)
+    print x, "=>", y
+
+  test("10.0.0.44")
+  test("010.000.000.044")
+  test("::1")
+  test("::")