Make countries to check for bridge blocking configurable

countries-that-block-bridges.py

 #!/usr/bin/env python3
 
-import collections
 import csv
 import datetime
+import json
 import statistics
 
 import pendulum
 from tqdm import tqdm
-import json
 
 from analyzer import OoniAnalyser
 
@@ -18,7 +17,6 @@ UNIT = 'days'  # hours, days, weeks
 INTERVAL = 1
 BRIDGE_CONTROL_THRESHOLD = 80  # only consider bridges with a control success rate higher than this
 
-COUNTRIES = BLOCKING_COUNTRIES + CONTROL_COUNTRIES
 FAILURES = [
     'generic_timeout_error',
     'connection_refused_error',
@@ -35,16 +33,25 @@ def main():
 class BridgeAnalyzer(OoniAnalyser):
     pre_report_init = False
     period = None
+    blocking_countries = []
+    countries = []
     bridge_data = {}
 
     def init_arg_parser(self):
         parser = super().init_arg_parser()
         parser.add_argument('-n', '--nicknames', dest='nicknames', action='store_true',
                             help='replace bridge addresses with nicknames')
+        parser.add_argument('-c', '--countries', nargs='+', metavar='COUNTRY_CODE',
+                            default=BLOCKING_COUNTRIES, help='list of country codes to consider')
         return parser
 
+    def analyze(self):
+        self.blocking_countries = [country.upper() for country in self.args.countries]
+        self.countries = self.blocking_countries + CONTROL_COUNTRIES
+        super().analyze()
+
     def use_country(self, cc):
-        return cc in COUNTRIES
+        return cc in self.countries
 
     @staticmethod
     def use_report(data):
@@ -108,12 +115,11 @@ class BridgeAnalyzer(OoniAnalyser):
                     data['countries'][cc]['success'] += country['success']
                     data['countries'][cc]['failure'] += country['failure']
 
-    @staticmethod
-    def get_field_names():
-        blocking_total = [cc + "#" for cc in BLOCKING_COUNTRIES]
+    def get_field_names(self):
+        blocking_total = [cc + "#" for cc in self.blocking_countries]
         control_total = [cc + "#" for cc in CONTROL_COUNTRIES]
         return ('bridge', 'type', 'total') + \
-            tuple(sorted(BLOCKING_COUNTRIES)) + \
+            tuple(sorted(self.blocking_countries)) + \
             tuple(['ctrl'] + sorted(CONTROL_COUNTRIES)) + \
             tuple(sorted(blocking_total)) + \
             tuple(sorted(control_total))
@@ -155,7 +161,7 @@ class BridgeAnalyzer(OoniAnalyser):
                 continue
             # remember success rates for blocking countries
             if output_dict[bridge]['ctrl'] >= BRIDGE_CONTROL_THRESHOLD:
-                for cc in BLOCKING_COUNTRIES:
+                for cc in self.blocking_countries:
                     bridge_type = output_dict[bridge]["type"]
                     if bridge_type == "B" or bridge_type == "TB":
                         if cc not in output_dict[bridge]: