Add Python wrappers for bind and connect as well as sendmsg.

Change-Id: Ie75ec2d56a58a74e1680eb01993797a48cec7c70

3 files changed, 31 insertions, 14 deletions

diff --git a/tests/net_test/sendmsg.py b/tests/net_test/csocket.py
index 422a16a1..4b268e92 100644
--- a/tests/net_test/sendmsg.py
+++ b/tests/net_test/csocket.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Python wrapper for sendmsg."""
+"""Python wrapper for C socket calls and data structures."""
 
 import ctypes
 import ctypes.util
@@ -23,7 +23,7 @@ import struct
 import cstruct
 
 
-# Data structures used by sendmsg.
+# Data structures.
 CMsgHdr = cstruct.Struct("cmsghdr", "@Lii", "len level type")
 Iovec = cstruct.Struct("iovec", "@LL", "base len")
 MsgHdr = cstruct.Struct("msghdr", "@LLLLLLi",
@@ -44,6 +44,12 @@ def PaddedLength(length):
   return CMSG_ALIGNTO * ((length / CMSG_ALIGNTO) + (length % CMSG_ALIGNTO != 0))
 
 
+def MaybeRaiseSocketError(ret):
+  if ret < 0:
+    errno = ctypes.get_errno()
+    raise socket.error(errno, os.strerror(errno))
+
+
 def Sockaddr(addr):
   if ":" in addr[0]:
     family = socket.AF_INET6
@@ -100,12 +106,25 @@ def _MakeMsgControl(optlist):
   return msg_control
 
 
+def Bind(s, to):
+  """Python wrapper for connect."""
+  ret = libc.bind(s.fileno(), to.CPointer(), len(to))
+  MaybeRaiseSocketError(ret)
+  return ret
+
+def Connect(s, to):
+  """Python wrapper for connect."""
+  ret = libc.connect(s.fileno(), to.CPointer(), len(to))
+  MaybeRaiseSocketError(ret)
+  return ret
+
+
 def Sendmsg(s, to, data, control, flags):
   """Python wrapper for sendmsg.
 
   Args:
     s: A Python socket object. Becomes sockfd.
-    to: A Python socket address tuple. Becomes msg->msg_name.
+    to: An address tuple, or a SockaddrIn[6] struct. Becomes msg->msg_name.
     data: A string, the data to write. Goes into msg->msg_iov.
     control: A list of cmsg options. Becomes msg->msg_control.
     flags: An integer. Becomes msg->msg_flags.
@@ -122,9 +141,10 @@ def Sendmsg(s, to, data, control, flags):
 
   # Convert the destination address into a struct sockaddr.
   if to:
-    name = Sockaddr(to)
-    msg_name = name.CPointer()
-    msg_namelen = len(name)
+    if isinstance(to, tuple):
+      to = Sockaddr(to)
+    msg_name = to.CPointer()
+    msg_namelen = len(to)
   else:
     msg_name = 0
     msg_namelen = 0
@@ -155,9 +175,6 @@ def Sendmsg(s, to, data, control, flags):
 
   # Call sendmsg.
   ret = libc.sendmsg(s.fileno(), msghdr, 0)
-
-  if ret < 0:
-    errno = ctypes.get_errno()
-    raise socket.error(errno, os.strerror(errno))
+  MaybeRaiseSocketError(ret)
 
   return ret
diff --git a/tests/net_test/multinetwork_base.py b/tests/net_test/multinetwork_base.py
index e3cda7a0..89402588 100644
--- a/tests/net_test/multinetwork_base.py
+++ b/tests/net_test/multinetwork_base.py
@@ -27,10 +27,10 @@ import struct
 
 from scapy import all as scapy
 
+import csocket
 import cstruct
 import iproute
 import net_test
-import sendmsg
 
 
 IFF_TUN = 1
@@ -447,7 +447,7 @@ class MultiNetworkBaseTest(net_test.NetworkTest):
           4: (net_test.SOL_IP, IP_PKTINFO),
           6: (net_test.SOL_IPV6, IPV6_PKTINFO)}[version]
       cmsgs.append((cmsg_level, cmsg_name, pktinfo))
-    sendmsg.Sendmsg(s, (dstaddr, dstport), payload, cmsgs, sendmsg.MSG_CONFIRM)
+    csocket.Sendmsg(s, (dstaddr, dstport), payload, cmsgs, csocket.MSG_CONFIRM)
 
   def ReceiveEtherPacketOn(self, netid, packet):
     posix.write(self.tuns[netid].fileno(), str(packet))
diff --git a/tests/net_test/srcaddr_selection_test.py b/tests/net_test/srcaddr_selection_test.py
index 323c7dab..eb09b7fb 100755
--- a/tests/net_test/srcaddr_selection_test.py
+++ b/tests/net_test/srcaddr_selection_test.py
@@ -23,11 +23,11 @@ import unittest
 
 from scapy import all as scapy
 
+import csocket
 import iproute
 import multinetwork_base
 import multinetwork_test
 import net_test
-import sendmsg
 
 # Setsockopt values.
 IPV6_ADDR_PREFERENCES = 72
@@ -88,7 +88,7 @@ class IPv6SourceAddressSelectionTest(multinetwork_base.MultiNetworkBaseTest):
     pktinfo = multinetwork_base.MakePktInfo(6, address, 0)
     cmsgs = [(net_test.SOL_IPV6, IPV6_PKTINFO, pktinfo)]
     s = self.BuildSocket(6, net_test.UDPSocket, netid, "mark")
-    return sendmsg.Sendmsg(s, (dest, 53), "Hello", cmsgs, 0)
+    return csocket.Sendmsg(s, (dest, 53), "Hello", cmsgs, 0)
 
   def assertAddressUsable(self, address, netid):
     self.BindToAddress(address)