Merge pull request #1357 from sarthurdev/geoip

firewall: T4299: Add support for GeoIP filtering

10 files changed, 332 insertions, 2 deletions

diff --git a/data/templates/firewall/nftables-geoip-update.j2 b/data/templates/firewall/nftables-geoip-update.j2
new file mode 100644
index 000000000..f9e61a274
--- /dev/null
+++ b/data/templates/firewall/nftables-geoip-update.j2
@@ -0,0 +1,33 @@
+#!/usr/sbin/nft -f
+
+{% if ipv4_sets is vyos_defined %}
+{%     for setname, ip_list in ipv4_sets.items() %}
+flush set ip filter {{ setname }}
+{%     endfor %}
+
+table ip filter {
+{%     for setname, ip_list in ipv4_sets.items() %}
+    set {{ setname }} {
+        type ipv4_addr
+        flags interval
+        elements = { {{ ','.join(ip_list) }} }
+    }
+{%     endfor %}
+}
+{% endif %}
+
+{% if ipv6_sets is vyos_defined %}
+{%     for setname, ip_list in ipv6_sets.items() %}
+flush set ip6 filter {{ setname }}
+{%     endfor %}
+
+table ip6 filter {
+{%     for setname, ip_list in ipv6_sets.items() %}
+    set {{ setname }} {
+        type ipv6_addr
+        flags interval
+        elements = { {{ ','.join(ip_list) }} }
+    }
+{%     endfor %}
+}
+{% endif %}
diff --git a/data/templates/firewall/nftables.j2 b/data/templates/firewall/nftables.j2
index 1f88ae40c..961b83301 100644
--- a/data/templates/firewall/nftables.j2
+++ b/data/templates/firewall/nftables.j2
@@ -60,6 +60,14 @@ table ip filter {
         flags dynamic
     }
 {%     endfor %}
+{%     if geoip_updated.name is vyos_defined %}
+{%         for setname in geoip_updated.name %}
+    set {{ setname }} {
+        type ipv4_addr
+        flags interval
+    }
+{%         endfor %}
+{%     endif %}
 {% endif %}
 {% if state_policy is vyos_defined %}
     chain VYOS_STATE_POLICY {
@@ -121,6 +129,14 @@ table ip6 filter {
         flags dynamic
     }
 {%     endfor %}
+{%     if geoip_updated.ipv6_name is vyos_defined %}
+{%         for setname in geoip_updated.ipv6_name %}
+    set {{ setname }} {
+        type ipv6_addr
+        flags interval
+    }
+{%         endfor %}
+{%     endif %}
 {% endif %}
 {% if state_policy is vyos_defined %}
     chain VYOS_STATE_POLICY6 {
diff --git a/interface-definitions/firewall.xml.in b/interface-definitions/firewall.xml.in
index 719088d03..2e9452dfd 100644
--- a/interface-definitions/firewall.xml.in
+++ b/interface-definitions/firewall.xml.in
@@ -366,6 +366,7 @@
                 </properties>
                 <children>
                   #include <include/firewall/address-ipv6.xml.i>
+                  #include <include/firewall/geoip.xml.i>
                   #include <include/firewall/source-destination-group-ipv6.xml.i>
                   #include <include/firewall/port.xml.i>
                 </children>
@@ -376,6 +377,7 @@
                 </properties>
                 <children>
                   #include <include/firewall/address-ipv6.xml.i>
+                  #include <include/firewall/geoip.xml.i>
                   #include <include/firewall/source-destination-group-ipv6.xml.i>
                   #include <include/firewall/port.xml.i>
                 </children>
@@ -552,6 +554,7 @@
                 </properties>
                 <children>
                   #include <include/firewall/address.xml.i>
+                  #include <include/firewall/geoip.xml.i>
                   #include <include/firewall/source-destination-group.xml.i>
                   #include <include/firewall/port.xml.i>
                 </children>
@@ -562,6 +565,7 @@
                 </properties>
                 <children>
                   #include <include/firewall/address.xml.i>
+                  #include <include/firewall/geoip.xml.i>
                   #include <include/firewall/source-destination-group.xml.i>
                   #include <include/firewall/port.xml.i>
                 </children>
diff --git a/interface-definitions/include/firewall/geoip.xml.i b/interface-definitions/include/firewall/geoip.xml.i
new file mode 100644
index 000000000..f6208f718
--- /dev/null
+++ b/interface-definitions/include/firewall/geoip.xml.i
@@ -0,0 +1,22 @@
+<!-- include start from firewall/geoip.xml.i -->
+<node name="geoip">
+  <properties>
+    <help>GeoIP options - Data provided by DB-IP.com</help>
+  </properties>
+  <children>
+    <leafNode name="country-code">
+      <properties>
+        <help>GeoIP country code</help>
+        <valueHelp>
+          <format>&lt;country&gt;</format>
+          <description>Country code (2 characters)</description>
+        </valueHelp>
+        <constraint>
+          <regex>^(ad|ae|af|ag|ai|al|am|ao|aq|ar|as|at|au|aw|ax|az|ba|bb|bd|be|bf|bg|bh|bi|bj|bl|bm|bn|bo|bq|br|bs|bt|bv|bw|by|bz|ca|cc|cd|cf|cg|ch|ci|ck|cl|cm|cn|co|cr|cu|cv|cw|cx|cy|cz|de|dj|dk|dm|do|dz|ec|ee|eg|eh|er|es|et|fi|fj|fk|fm|fo|fr|ga|gb|gd|ge|gf|gg|gh|gi|gl|gm|gn|gp|gq|gr|gs|gt|gu|gw|gy|hk|hm|hn|hr|ht|hu|id|ie|il|im|in|io|iq|ir|is|it|je|jm|jo|jp|ke|kg|kh|ki|km|kn|kp|kr|kw|ky|kz|la|lb|lc|li|lk|lr|ls|lt|lu|lv|ly|ma|mc|md|me|mf|mg|mh|mk|ml|mm|mn|mo|mp|mq|mr|ms|mt|mu|mv|mw|mx|my|mz|na|nc|ne|nf|ng|ni|nl|no|np|nr|nu|nz|om|pa|pe|pf|pg|ph|pk|pl|pm|pn|pr|ps|pt|pw|py|qa|re|ro|rs|ru|rw|sa|sb|sc|sd|se|sg|sh|si|sj|sk|sl|sm|sn|so|sr|ss|st|sv|sx|sy|sz|tc|td|tf|tg|th|tj|tk|tl|tm|tn|to|tr|tt|tv|tw|tz|ua|ug|um|us|uy|uz|va|vc|ve|vg|vi|vn|vu|wf|ws|ye|yt|za|zm|zw)$</regex>
+        </constraint>
+        <multi />
+      </properties>
+    </leafNode>
+  </children>
+</node>
+<!-- include end -->
diff --git a/op-mode-definitions/geoip.xml.in b/op-mode-definitions/geoip.xml.in
new file mode 100644
index 000000000..c1b6e87b9
--- /dev/null
+++ b/op-mode-definitions/geoip.xml.in
@@ -0,0 +1,13 @@
+<?xml version="1.0"?>
+<interfaceDefinition>
+  <node name="update">
+    <children>
+      <leafNode name="geoip">
+        <properties>
+          <help>Update GeoIP database and firewall sets</help>
+        </properties>
+        <command>sudo ${vyos_libexec_dir}/geoip-update.py --force</command>
+      </leafNode>
+    </children>
+  </node>
+</interfaceDefinition>
diff --git a/python/vyos/firewall.py b/python/vyos/firewall.py
index 355ec44b0..a61d0a9f8 100644
--- a/python/vyos/firewall.py
+++ b/python/vyos/firewall.py
@@ -14,11 +14,22 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import csv
+import gzip
+import os
 import re
 
+from pathlib import Path
+from time import strftime
+
+from vyos.remote import download
+from vyos.template import is_ipv4
+from vyos.template import render
 from vyos.util import call
 from vyos.util import cmd
 from vyos.util import dict_search_args
+from vyos.util import dict_search_recursive
+from vyos.util import run
 
 
 # Functions for firewall group domain-groups
@@ -139,6 +150,9 @@ def parse_rule(rule_conf, fw_name, rule_id, ip_name):
                 if suffix[0] == '!':
                     suffix = f'!= {suffix[1:]}'
                 output.append(f'{ip_name} {prefix}addr {suffix}')
+            
+            if dict_search_args(side_conf, 'geoip', 'country_code'):
+                output.append(f'{ip_name} {prefix}addr @GEOIP_CC_{fw_name}_{rule_id}')
 
             if 'mac_address' in side_conf:
                 suffix = side_conf["mac_address"]
@@ -338,3 +352,127 @@ def parse_policy_set(set_conf, def_suffix):
         mss = set_conf['tcp_mss']
         out.append(f'tcp option maxseg size set {mss}')
     return " ".join(out)
+
+# GeoIP
+
+nftables_geoip_conf = '/run/nftables-geoip.conf'
+geoip_database = '/usr/share/vyos-geoip/dbip-country-lite.csv.gz'
+geoip_lock_file = '/run/vyos-geoip.lock'
+
+def geoip_load_data(codes=[]):
+    data = None
+
+    if not os.path.exists(geoip_database):
+        return []
+
+    try:
+        with gzip.open(geoip_database, mode='rt') as csv_fh:
+            reader = csv.reader(csv_fh)
+            out = []
+            for start, end, code in reader:
+                if code.lower() in codes:
+                    out.append([start, end, code.lower()])
+            return out
+    except:
+        print('Error: Failed to open GeoIP database')
+    return []
+
+def geoip_download_data():
+    url = 'https://download.db-ip.com/free/dbip-country-lite-{}.csv.gz'.format(strftime("%Y-%m"))
+    try:
+        dirname = os.path.dirname(geoip_database)
+        if not os.path.exists(dirname):
+            os.mkdir(dirname)
+
+        download(geoip_database, url)
+        print("Downloaded GeoIP database")
+        return True
+    except:
+        print("Error: Failed to download GeoIP database")
+    return False
+
+class GeoIPLock(object):
+    def __init__(self, file):
+        self.file = file
+
+    def __enter__(self):
+        if os.path.exists(self.file):
+            return False
+
+        Path(self.file).touch()
+        return True
+
+    def __exit__(self, exc_type, exc_value, tb):
+        os.unlink(self.file)
+
+def geoip_update(firewall, force=False):
+    with GeoIPLock(geoip_lock_file) as lock:
+        if not lock:
+            print("Script is already running")
+            return False
+
+        if not firewall:
+            print("Firewall is not configured")
+            return True
+
+        if not os.path.exists(geoip_database):
+            if not geoip_download_data():
+                return False
+        elif force:
+            geoip_download_data()
+
+        ipv4_codes = {}
+        ipv6_codes = {}
+
+        ipv4_sets = {}
+        ipv6_sets = {}
+
+        # Map country codes to set names
+        for codes, path in dict_search_recursive(firewall, 'country_code'):
+            if path[0] == 'name':
+                set_name = f'GEOIP_CC_{path[1]}_{path[3]}'
+                ipv4_sets[set_name] = []
+                for code in codes:
+                    if code not in ipv4_codes:
+                        ipv4_codes[code] = [set_name]
+                    else:
+                        ipv4_codes[code].append(set_n)
+            elif path[0] == 'ipv6_name':
+                set_name = f'GEOIP_CC_{path[1]}_{path[3]}'
+                ipv6_sets[set_name] = []
+                for code in codes:
+                    if code not in ipv6_codes:
+                        ipv6_codes[code] = [set_name]
+                    else:
+                        ipv6_codes[code].append(set_name)
+
+        if not ipv4_codes and not ipv6_codes:
+            if force:
+                print("GeoIP not in use by firewall")
+            return True
+
+        geoip_data = geoip_load_data([*ipv4_codes, *ipv6_codes])
+
+        # Iterate IP blocks to assign to sets
+        for start, end, code in geoip_data:
+            ipv4 = is_ipv4(start)
+            if code in ipv4_codes and ipv4:
+                ip_range = f'{start}-{end}' if start != end else start
+                for setname in ipv4_codes[code]:
+                    ipv4_sets[setname].append(ip_range)
+            if code in ipv6_codes and not ipv4:
+                ip_range = f'{start}-{end}' if start != end else start
+                for setname in ipv6_codes[code]:
+                    ipv6_sets[setname].append(ip_range)
+
+        render(nftables_geoip_conf, 'firewall/nftables-geoip-update.j2', {
+            'ipv4_sets': ipv4_sets,
+            'ipv6_sets': ipv6_sets
+        })
+
+        result = run(f'nft -f {nftables_geoip_conf}')
+        if result != 0:
+            print('Error: GeoIP failed to update firewall')
+            return False
+
+        return True
diff --git a/src/conf_mode/firewall.py b/src/conf_mode/firewall.py
index 792e17b85..46b8add59 100755
--- a/src/conf_mode/firewall.py
+++ b/src/conf_mode/firewall.py
@@ -26,6 +26,7 @@ from vyos.config import Config
 from vyos.configdict import dict_merge
 from vyos.configdict import node_changed
 from vyos.configdiff import get_config_diff, Diff
+from vyos.firewall import geoip_update
 from vyos.firewall import get_ips_domains_dict
 from vyos.firewall import nft_add_set_elements
 from vyos.firewall import nft_flush_set
@@ -35,6 +36,7 @@ from vyos.template import render
 from vyos.util import call
 from vyos.util import cmd
 from vyos.util import dict_search_args
+from vyos.util import dict_search_recursive
 from vyos.util import process_named_running
 from vyos.util import run
 from vyos.xml import defaults
@@ -150,6 +152,38 @@ def get_firewall_zones(conf):
 
     return {'name': used_v4, 'ipv6_name': used_v6}
 
+def geoip_updated(conf, firewall):
+    diff = get_config_diff(conf)
+    node_diff = diff.get_child_nodes_diff(['firewall'], expand_nodes=Diff.DELETE, recursive=True)
+
+    out = {
+        'name': [],
+        'ipv6_name': [],
+        'deleted_name': [],
+        'deleted_ipv6_name': []
+    }
+    updated = False
+
+    for key, path in dict_search_recursive(firewall, 'geoip'):
+        if path[0] == 'name':
+            out['name'].append(f'GEOIP_CC_{path[1]}_{path[3]}')
+        elif path[0] == 'ipv6_name':
+            out['ipv6_name'].append(f'GEOIP_CC_{path[1]}_{path[3]}')
+        updated = True
+
+    if 'delete' in node_diff:
+        for key, path in dict_search_recursive(node_diff['delete'], 'geoip'):
+            if path[0] == 'name':
+                out['deleted_name'].append(f'GEOIP_CC_{path[1]}_{path[3]}')
+            elif path[0] == 'ipv6-name':
+                out['deleted_ipv6_name'].append(f'GEOIP_CC_{path[1]}_{path[3]}')
+            updated = True
+
+    if updated:
+        return out
+
+    return False
+
 def get_config(config=None):
     if config:
         conf = config
@@ -174,6 +208,8 @@ def get_config(config=None):
                                         key_mangling=('-', '_'), get_first_key=True,
                                         no_tag_node_value_mangle=True)
 
+    firewall['geoip_updated'] = geoip_updated(conf, firewall)
+
     return firewall
 
 def verify_rule(firewall, rule_conf, ipv6):
@@ -219,6 +255,16 @@ def verify_rule(firewall, rule_conf, ipv6):
         if side in rule_conf:
             side_conf = rule_conf[side]
 
+            if dict_search_args(side_conf, 'geoip', 'country_code'):
+                if 'address' in side_conf:
+                    raise ConfigError('Address and GeoIP cannot both be defined')
+
+                if dict_search_args(side_conf, 'group', 'address_group'):
+                    raise ConfigError('Address-group and GeoIP cannot both be defined')
+
+                if dict_search_args(side_conf, 'group', 'network_group'):
+                    raise ConfigError('Network-group and GeoIP cannot both be defined')
+
             if 'group' in side_conf:
                 if {'address_group', 'network_group'} <= set(side_conf['group']):
                     raise ConfigError('Only one address-group or network-group can be specified')
@@ -322,8 +368,13 @@ def cleanup_commands(firewall):
     commands = []
     commands_end = []
     for table in ['ip filter', 'ip6 filter']:
+        geoip_list = []
+        if firewall['geoip_updated']:
+            geoip_key = 'deleted_ipv6_name' if table == 'ip6 filter' else 'deleted_name'
+            geoip_list = dict_search_args(firewall, 'geoip_updated', geoip_key) or []
+
         state_chain = 'VYOS_STATE_POLICY' if table == 'ip filter' else 'VYOS_STATE_POLICY6'
-        json_str = cmd(f'nft -j list table {table}')
+        json_str = cmd(f'nft -t -j list table {table}')
         obj = loads(json_str)
         if 'nftables' not in obj:
             continue
@@ -353,6 +404,8 @@ def cleanup_commands(firewall):
                             commands.append(f'delete rule {table} {chain} handle {handle}')
             elif 'set' in item:
                 set_name = item['set']['name']
+                if set_name.startswith('GEOIP_CC_') and set_name not in geoip_list:
+                    continue
                 commands_end.append(f'delete set {table} {set_name}')
     return commands + commands_end
 
@@ -476,6 +529,12 @@ def apply(firewall):
     if firewall['policy_resync']:
         resync_policy_route()
 
+    if firewall['geoip_updated']:
+        # Call helper script to Update set contents
+        if 'name' in firewall['geoip_updated'] or 'ipv6_name' in firewall['geoip_updated']:
+            print('Updating GeoIP. Please wait...')
+            geoip_update(firewall)
+
     post_apply_trap(firewall)
 
     return None
diff --git a/src/conf_mode/zone_policy.py b/src/conf_mode/zone_policy.py
index 070a4deea..a52c52706 100755
--- a/src/conf_mode/zone_policy.py
+++ b/src/conf_mode/zone_policy.py
@@ -155,7 +155,7 @@ def get_local_from(zone_policy, local_zone_name):
 def cleanup_commands():
     commands = []
     for table in ['ip filter', 'ip6 filter']:
-        json_str = cmd(f'nft -j list table {table}')
+        json_str = cmd(f'nft -t -j list table {table}')
         obj = loads(json_str)
         if 'nftables' not in obj:
             continue
diff --git a/src/etc/cron.d/vyos-geoip b/src/etc/cron.d/vyos-geoip
new file mode 100644
index 000000000..9bb38a850
--- /dev/null
+++ b/src/etc/cron.d/vyos-geoip
@@ -0,0 +1 @@
+30 4 * * 1 root sg vyattacfg "/usr/libexec/vyos/geoip-update.py --force" >/tmp/geoip-update.log 2>&1
diff --git a/src/helpers/geoip-update.py b/src/helpers/geoip-update.py
new file mode 100755
index 000000000..34accf2cc
--- /dev/null
+++ b/src/helpers/geoip-update.py
@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+#
+# Copyright (C) 2021 VyOS maintainers and contributors
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 2 or later as
+# published by the Free Software Foundation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import argparse
+import sys
+
+from vyos.configquery import ConfigTreeQuery
+from vyos.firewall import geoip_update
+
+def get_config(config=None):
+    if config:
+        conf = config
+    else:
+        conf = ConfigTreeQuery()
+    base = ['firewall']
+
+    if not conf.exists(base):
+        return None
+
+    return conf.get_config_dict(base, key_mangling=('-', '_'), get_first_key=True,
+                                    no_tag_node_value_mangle=True)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--force", help="Force update", action="store_true")
+    args = parser.parse_args()
+
+    firewall = get_config()
+
+    if not geoip_update(firewall, force=args.force):
+        sys.exit(1)