Pass group to firewall

5 files changed, 29 insertions, 21 deletions

diff --git a/sshuttle/client.py b/sshuttle/client.py
index 25b3440..5d1d990 100644
--- a/sshuttle/client.py
+++ b/sshuttle/client.py
@@ -319,7 +319,7 @@ class FirewallClient:
 
     def setup(self, subnets_include, subnets_exclude, nslist,
               redirectport_v6, redirectport_v4, dnsport_v6, dnsport_v4, udp,
-              user, tmark):
+              user, group, tmark):
         self.subnets_include = subnets_include
         self.subnets_exclude = subnets_exclude
         self.nslist = nslist
@@ -329,6 +329,7 @@ class FirewallClient:
         self.dnsport_v4 = dnsport_v4
         self.udp = udp
         self.user = user
+        self.group = group
         self.tmark = tmark
 
     def check(self):
@@ -367,9 +368,14 @@ class FirewallClient:
             user = bytes(self.user, 'utf-8')
         else:
             user = b'%d' % self.user
-
-        self.pfile.write(b'GO %d %s %s %d\n' %
-                         (udp, user, bytes(self.tmark, 'ascii'), os.getpid()))
+        if self.group is None:
+            group = b'-'
+        elif isinstance(self.group, str):
+            group =  bytes(self.group, 'utf-8')
+        else:
+            group = b'%d' % self.group
+        self.pfile.write(b'GO %d %s %s %s %d\n' %
+                         (udp, user, group, bytes(self.tmark, 'ascii'), os.getpid()))
         self.pfile.flush()
 
         line = self.pfile.readline()
diff --git a/sshuttle/firewall.py b/sshuttle/firewall.py
index af71fe7..60662b9 100644
--- a/sshuttle/firewall.py
+++ b/sshuttle/firewall.py
@@ -270,13 +270,15 @@ def main(method_name, syslog):
 
     _, _, args = line.partition(" ")
     global sshuttle_pid
-    udp, user, tmark, sshuttle_pid = args.strip().split(" ", 3)
+    udp, user, group, tmark, sshuttle_pid = args.strip().split(" ", 4)
     udp = bool(int(udp))
     sshuttle_pid = int(sshuttle_pid)
     if user == '-':
         user = None
-    debug2('Got udp: %r, user: %r, tmark: %s, sshuttle_pid: %d' %
-           (udp, user, tmark, sshuttle_pid))
+    if group == '-':
+        group = None
+    debug2('Got udp: %r, user: %r, group: %r, tmark: %s, sshuttle_pid: %d' %
+           (udp, user, group, tmark, sshuttle_pid))
 
     subnets_v6 = [i for i in subnets if i[0] == socket.AF_INET6]
     nslist_v6 = [i for i in nslist if i[0] == socket.AF_INET6]
@@ -291,14 +293,14 @@ def main(method_name, syslog):
             method.setup_firewall(
                 port_v6, dnsport_v6, nslist_v6,
                 socket.AF_INET6, subnets_v6, udp,
-                user, tmark)
+                user, group, tmark)
 
         if subnets_v4 or nslist_v4:
             debug2('setting up IPv4.')
             method.setup_firewall(
                 port_v4, dnsport_v4, nslist_v4,
                 socket.AF_INET, subnets_v4, udp,
-                user, tmark)
+                user, group, tmark)
 
         flush_systemd_dns_cache()
         stdout.write('STARTED\n')
@@ -334,7 +336,7 @@ def main(method_name, syslog):
         try:
             if subnets_v6 or nslist_v6:
                 debug2('undoing IPv6 changes.')
-                method.restore_firewall(port_v6, socket.AF_INET6, udp, user)
+                method.restore_firewall(port_v6, socket.AF_INET6, udp, user, group)
         except Exception:
             try:
                 debug1("Error trying to undo IPv6 firewall.")
@@ -345,7 +347,7 @@ def main(method_name, syslog):
         try:
             if subnets_v4 or nslist_v4:
                 debug2('undoing IPv4 changes.')
-                method.restore_firewall(port_v4, socket.AF_INET, udp, user)
+                method.restore_firewall(port_v4, socket.AF_INET, udp, user, group)
         except Exception:
             try:
                 debug1("Error trying to undo IPv4 firewall.")
diff --git a/sshuttle/methods/__init__.py b/sshuttle/methods/__init__.py
index 4a1abe6..962529b 100644
--- a/sshuttle/methods/__init__.py
+++ b/sshuttle/methods/__init__.py
@@ -90,10 +90,10 @@ class BaseMethod(object):
                     (key, self.name))
 
     def setup_firewall(self, port, dnsport, nslist, family, subnets, udp,
-                       user, tmark):
+                       user, group, tmark):
         raise NotImplementedError()
 
-    def restore_firewall(self, port, family, udp, user):
+    def restore_firewall(self, port, family, udp, user, group):
         raise NotImplementedError()
 
     @staticmethod
diff --git a/sshuttle/methods/nat.py b/sshuttle/methods/nat.py
index 1254e55..bd878b1 100644
--- a/sshuttle/methods/nat.py
+++ b/sshuttle/methods/nat.py
@@ -31,17 +31,17 @@ class Method(BaseMethod):
         chain = 'sshuttle-%s' % port
 
         # basic cleanup/setup of chains
-        self.restore_firewall(port, family, udp, user)
+        self.restore_firewall(port, family, udp, user, group)
 
         _ipt('-N', chain)
         _ipt('-F', chain)
         if user is not None or group is not None:
             margs = ['-I', 'OUTPUT', '1', '-m', 'owner']
             if user is not None:
-                margs.append('--uid-owner', str(user))
+                margs += ['--uid-owner', str(user)]
             if group is not None:
-                margs.append('--gid-owner', str(group))
-            margs = args.append('-j', 'MARK', '--set-mark', str(port))
+                margs += ['--gid-owner', str(group)]
+            margs += ['-j', 'MARK', '--set-mark', str(port)]
             nonfatal(_ipm, *margs)
             args = '-m', 'mark', '--mark', str(port), '-j', chain
         else:
@@ -104,10 +104,10 @@ class Method(BaseMethod):
             if user is not None or group is not None:
                 margs = ['-D', 'OUTPUT', '-m', 'owner']
                 if user is not None:
-                    margs.append('--uid-owner', str(user))
+                    margs += ['--uid-owner', str(user)]
                 if group is not None:
-                    margs.append('--gid-owner', str(group))
-                margs = args.append('-j', 'MARK', '--set-mark', str(port))
+                    margs += ['--gid-owner', str(group)]
+                margs += ['-j', 'MARK', '--set-mark', str(port)]
                 nonfatal(_ipm, *margs)
 
                 args = '-m', 'mark', '--mark', str(port), '-j', chain
diff --git a/sshuttle/methods/tproxy.py b/sshuttle/methods/tproxy.py
index e12943c..b3d5fca 100644
--- a/sshuttle/methods/tproxy.py
+++ b/sshuttle/methods/tproxy.py
@@ -134,7 +134,7 @@ class Method(BaseMethod):
         divert_chain = 'sshuttle-d-%s' % port
 
         # basic cleanup/setup of chains
-        self.restore_firewall(port, family, udp, user)
+        self.restore_firewall(port, family, udp, user, group)
 
         _ipt('-N', mark_chain)
         _ipt('-F', mark_chain)