Merge pull request #4 from seanzxx/yosemite_support

Yosemite support

3 files changed, 125 insertions, 3 deletions

diff --git a/src/client.py b/src/client.py
index 6b7a293..f919b37 100644
--- a/src/client.py
+++ b/src/client.py
@@ -12,6 +12,8 @@ import ssyslog
 import sys
 from ssnet import SockWrapper, Handler, Proxy, Mux, MuxWrapper
 from helpers import log, debug1, debug2, debug3, Fatal, islocal
+from fcntl import ioctl
+from ctypes import c_char, c_uint8, c_uint16, c_uint32, Union, Structure, sizeof, addressof, memmove
 
 recvmsg = None
 try:
@@ -185,6 +187,79 @@ def daemon_cleanup():
             raise
 
 
+class pf_state_xport(Union):
+    _fields_ = [("port", c_uint16),
+                ("call_id", c_uint16),
+                ("spi", c_uint32)]
+
+class pf_addr(Structure):
+    class _pfa(Union):
+        _fields_ = [("v4",            c_uint32),      # struct in_addr
+                    ("v6",            c_uint32 * 4),  # struct in6_addr
+                    ("addr8",         c_uint8 * 16),
+                    ("addr16",        c_uint16 * 8),
+                    ("addr32",        c_uint32 * 4)]
+
+    _fields_ = [("pfa",               _pfa)]
+    _anonymous_ = ("pfa",)
+
+class pfioc_natlook(Structure):
+    _fields_ = [("saddr", pf_addr),
+                ("daddr", pf_addr),
+                ("rsaddr", pf_addr),
+                ("rdaddr", pf_addr),
+                ("sxport", pf_state_xport),
+                ("dxport", pf_state_xport),
+                ("rsxport", pf_state_xport),
+                ("rdxport", pf_state_xport),
+                ("af", c_uint8),                      # sa_family_t  
+                ("proto", c_uint8),
+                ("proto_variant", c_uint8),
+                ("direction", c_uint8)]
+
+DIOCNATLOOK = ((0x40000000L | 0x80000000L) | ((sizeof(pfioc_natlook) & 0x1fff) << 16) | ((ord('D')) << 8) | (23))
+PF_OUT = 2
+
+_pf_fd = None
+
+def pf_dst(sock):
+    global _pf_fd
+    try:
+        peer = sock.getpeername()
+        proxy = sock.getsockname()
+
+        pnl = pfioc_natlook()
+        pnl.proto = socket.IPPROTO_TCP
+        pnl.direction = PF_OUT
+        if sock.family == socket.AF_INET:
+            pnl.af = socket.AF_INET
+            memmove(addressof(pnl.saddr), socket.inet_pton(socket.AF_INET, peer[0]), 4)
+            pnl.sxport.port = socket.htons(peer[1])
+            memmove(addressof(pnl.daddr), socket.inet_pton(socket.AF_INET, proxy[0]), 4)
+            pnl.dxport.port = socket.htons(proxy[1])
+        elif sock.family == socket.AF_INET6:
+            pnl.af = socket.AF_INET6
+            memmove(addressof(pnl.saddr), socket.inet_pton(socket.AF_INET6, peer[0]), 16)
+            pnl.sxport.port = socket.htons(peer[1])
+            memmove(addressof(pnl.daddr), socket.inet_pton(socket.AF_INET6, proxy[0]), 16)
+            pnl.dxport.port = socket.htons(proxy[1])
+
+        if _pf_fd == None:
+            _pf_fd = open('/dev/pf', 'r')
+
+        ioctl(_pf_fd, DIOCNATLOOK, (c_char * sizeof(pnl)).from_address(addressof(pnl)))
+
+        if pnl.af == socket.AF_INET:
+            ip = socket.inet_ntop(socket.AF_INET, (c_char * 4).from_address(addressof(pnl.rdaddr)))
+        elif pnl.af == socket.AF_INET6:
+            ip = socket.inet_ntop(socket.AF_INET6, (c_char * 16).from_address(addressof(pnl.rdaddr)))
+        port = socket.ntohs(pnl.rdxport.port)
+        return (ip, port)
+    except IOError, e:
+        return sock.getsockname()
+        raise
+
+
 def original_dst(sock):
     try:
         SO_ORIGINAL_DST = 80
@@ -381,6 +456,8 @@ def onaccept_tcp(listener, method, mux, handlers):
             raise
     if method == "tproxy":
         dstip = sock.getsockname()
+    elif method == "pf":
+        dstip = pf_dst(sock)
     else:
         dstip = original_dst(sock)
     debug1('Accept TCP: %s:%r -> %s:%r.\n' % (srcip[0], srcip[1],
diff --git a/src/firewall.py b/src/firewall.py
index 381913d..b255ee0 100644
--- a/src/firewall.py
+++ b/src/firewall.py
@@ -463,6 +463,47 @@ def do_ipfw(port, dnsport, family, subnets, udp):
     return do_wait
 
 
+def pfctl(*args):
+    argv = ['pfctl'] + list(args)
+    debug1('>> %s\n' % ' '.join(argv))
+    rv = ssubprocess.Popen(argv, stderr=ssubprocess.PIPE).wait()
+    if rv:
+        raise Fatal('%r returned %d' % (argv, rv))
+
+
+def do_pf(port, dnsport, family, subnets, udp):
+    tables = []
+    translating_rules = []
+    filtering_rules = []
+
+    if subnets:
+        include_subnets = filter(lambda s:not s[2], sorted(subnets, reverse=True))
+        if include_subnets:
+            tables.append('table <include_subnets> {%s}' % ','.join(["%s/%s" % (n[3], n[1]) for n in include_subnets]))
+            translating_rules.append('rdr pass on lo0 proto tcp to <include_subnets> -> 127.0.0.1 port %r' % port)
+            filtering_rules.append('pass out route-to lo0 inet proto tcp to <include_subnets> keep state')
+
+        exclude_subnets = filter(lambda s:s[2], sorted(subnets, reverse=True))
+        if exclude_subnets:
+            tables.append('table <exclude_subnets> {%s}' % ','.join(["%s/%s" % (n[3], n[1]) for n in exclude_subnets]))
+            filtering_rules.append('pass out route-to lo0 inet proto tcp to <exclude_subnets> keep state')
+
+        if dnsport:
+            nslist = resolvconf_nameservers()
+            tables.append('table <dns_servers> {%s}' % ','.join([ns[1] for ns in nslist]))
+            translating_rules.append('rdr pass on lo0 proto udp to <dns_servers> port 53 -> 127.0.0.1 port %r' % dnsport)
+            filtering_rules.append('pass out route-to lo0 inet proto udp to <dns_servers> port 53 keep state')
+
+        pf_config_file = '/etc/pf-sshuttle.conf'
+        with open(pf_config_file, 'w+') as f:
+            f.write('\n'.join(tables + translating_rules + filtering_rules) + '\n')
+
+        pfctl('-Ef', pf_config_file)
+        os.remove(pf_config_file)
+    else:
+        pfctl('-dF', 'all')
+
+
 def program_exists(name):
     paths = (os.getenv('PATH') or os.defpath).split(os.pathsep)
     for p in paths:
@@ -541,8 +582,10 @@ def main(port_v6, port_v4, dnsport_v6, dnsport_v4, method, udp, syslog):
             method = "ipfw"
         elif program_exists('iptables'):
             method = "nat"
+        elif program_exists('pfctl'):
+            method = "pf"
         else:
-            raise Fatal("can't find either ipfw or iptables; check your PATH")
+            raise Fatal("can't find either ipfw, iptables or pfctl; check your PATH")
 
     if method == "nat":
         do_it = do_iptables_nat
@@ -550,6 +593,8 @@ def main(port_v6, port_v4, dnsport_v6, dnsport_v4, method, udp, syslog):
         do_it = do_iptables_tproxy
     elif method == "ipfw":
         do_it = do_ipfw
+    elif method == "pf":
+        do_it = do_pf
     else:
         raise Exception('Unknown method "%s"' % method)
 
diff --git a/src/main.py b/src/main.py
index 98bac05..fe8275c 100644
--- a/src/main.py
+++ b/src/main.py
@@ -116,7 +116,7 @@ l,listen=  transproxy to this ip address and port number
 H,auto-hosts scan for remote hostnames and update local /etc/hosts
 N,auto-nets  automatically determine subnets to route
 dns        capture local DNS requests and forward to the remote DNS server
-method=    auto, nat, tproxy, or ipfw
+method=    auto, nat, tproxy, pf or ipfw
 python=    path to python interpreter on the remote server
 r,remote=  ssh hostname (and optional username) of remote sshuttle server
 x,exclude= exclude this subnet (can be used more than once)
@@ -183,7 +183,7 @@ try:
             includes = parse_subnet_file(opt.subnets)
         if not opt.method:
             method = "auto"
-        elif opt.method in ["auto", "nat", "tproxy", "ipfw"]:
+        elif opt.method in ["auto", "nat", "tproxy", "ipfw", "pf"]:
             method = opt.method
         else:
             o.fatal("method %s not supported" % opt.method)