arc.py: Add block_outbound arg in arc_setup()

The patch adds one more argument to block outbound network
traffic if required during a test. Also, we have to relax
localhost-bound connections as some tests require local
client-server connections.

BUG=chromium:705863
TEST=test_that cheets_BlockOutboundNetworkTest DUT_IP

Change-Id: Ib0cb892c9409e500b3307886a0030ee256de904f
Signed-off-by: Chung-yih Wang <cywang@google.com>
Reviewed-on: https://chromium-review.googlesource.com/461621
Reviewed-by: Ilja H. Friedel <ihf@chromium.org>
Reviewed-by: Ben Cheng <bccheng@chromium.org>
diff --git a/client/common_lib/cros/arc.py b/client/common_lib/cros/arc.py
index 2108c77..c955cdc 100644
--- a/client/common_lib/cros/arc.py
+++ b/client/common_lib/cros/arc.py
@@ -484,7 +484,8 @@
                     self._chrome.close()
 
     def arc_setup(self, dep_package=None, apks=None, full_pkg_names=[],
-                  uiautomator=False, email_id=None, password=None):
+                  uiautomator=False, email_id=None, password=None,
+                  block_outbound=False):
         """ARC test setup: Setup dependencies and install apks.
 
         This function disables package verification and enables non-market
@@ -500,6 +501,7 @@
         @param email_id: email id to be attached to the android. Only used
                          when  account_util is set to true.
         @param password: password related to the email_id.
+        @param block_outbound: block outbound network traffic during a test.
         """
         if not self.initialized:
             logging.info('Skipping ARC setup: not initialized')
@@ -557,6 +559,8 @@
         if self.uiautomator:
             path = os.path.join(self.autodir, 'deps', self._PKG_UIAUTOMATOR)
             sys.path.append(path)
+        if block_outbound:
+            self.block_outbound()
 
     def _stop_logcat(self):
         """Stop the adb logcat process gracefully."""
@@ -609,13 +613,12 @@
         """ Blocks the connection from the container to outer network.
 
             The iptables settings accept only 192.168.254.2 port 5555 (adb) and
-            localhost port 9008 (uiautomator)
+            all local connections, e.g. uiautomator.
         """
         logging.info('Blocking outbound connection')
         _android_shell('iptables -I OUTPUT -j REJECT')
         _android_shell('iptables -I OUTPUT -p tcp -s 192.168.254.2 --sport 5555 -j ACCEPT')
-        _android_shell('iptables -I OUTPUT -p tcp -d localhost --dport 9008 -j ACCEPT')
-        _android_shell('iptables -I OUTPUT -p tcp -s localhost --sport 9008 -j ACCEPT')
+        _android_shell('iptables -I OUTPUT -p tcp -d localhost -j ACCEPT')
 
 
     def unblock_outbound(self):
@@ -626,7 +629,6 @@
             unblock the outbound connections during the test if needed.
         """
         logging.info('Unblocking outbound connection')
-        _android_shell('iptables -D OUTPUT -p tcp -s localhost --sport 9008 -j ACCEPT')
-        _android_shell('iptables -D OUTPUT -p tcp -d localhost --dport 9008 -j ACCEPT')
+        _android_shell('iptables -D OUTPUT -p tcp -d localhost -j ACCEPT')
         _android_shell('iptables -D OUTPUT -p tcp -s 192.168.254.2 --sport 5555 -j ACCEPT')
         _android_shell('iptables -D OUTPUT -j REJECT')