Merge pull request #36 from chr0mag/sh0rch

Incorporate Sh0rch's enhancements
chr0mag · Nov 4, 2024 · 3b3c917 · 3b3c917
2 parents a44b725 + 5c94257
commit 3b3c917
Show file tree

Hide file tree

Showing 11 changed files with 173 additions and 111 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -25,7 +25,8 @@ jobs:
       - name: Generate geoipsets.conf
         working-directory: ./scripts
         env:
-          MAXMIND_KEY: ${{ secrets.MAXMIND_KEY }}
+          MAXMIND_ACCT_ID: ${{ secrets.MAXMIND_ACCT_ID }}
+          MAXMIND_NEW_KEY: ${{ secrets.MAXMIND_NEW_KEY }}
         run: |
           bash generate_geoipsets_conf.sh
       - name: Build geoipsets

diff --git a/.gitignore b/.gitignore
@@ -1,6 +1,8 @@
-.idea
+**/.vscode
+**/.idea
 *.iml
 **/.pytest_cache
+**/venv
 python/dist
 python/geoipsets.egg-info
 python/build

diff --git a/bash/bcs.conf b/bash/bcs.conf
@@ -1 +1,9 @@
-LICENSE_KEY=YOUR_LICENSE_KEY
+#Build GeoIP Sets fo netfilter 
+#Default configuration:
+#IPTABLES=no
+#NFTABLES=yes
+#IPv4=yes
+#IPv6=yes
+
+#Don't forget to set your MaxMind credentials!
+LICENSE_KEY=YOUR_ACCOUNT_ID:YOUR_LICENSE_KEY
diff --git a/bash/build-country-sets.sh b/bash/build-country-sets.sh
@@ -17,26 +17,27 @@ function error() {
 
 # ensure the programs needed to execute are available
 function check_progs() {
-  local PROGS="awk sed curl unzip md5sum cat mktemp"
+  local PROGS="awk sed curl unzip sha256sum cat mktemp"
   which ${PROGS} > /dev/null 2>&1 || error "Searching PATH fails to find executables among: ${PROGS}"
 }
 
 # retrieve latest MaxMind GeoLite2 IP country database and checksum
-# CSV URL: https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-Country-CSV&license_key=LICENSE_KEY&suffix=zip
-# MD5 URL: https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-Country-CSV&license_key=LICENSE_KEY&suffix=zip.md5
+# CSV URL: https://download.maxmind.com/geoip/databases/GeoLite2-Country-CSV/download?suffix=zip
+# SHA256 URL: https://download.maxmind.com/geoip/databases/GeoLite2-Country-CSV/download?suffix=zip.sha256
 function download_geolite2_data() {
   local ZIPPED_FILE="GeoLite2-Country-CSV.zip"
-  local MD5_FILE="${ZIPPED_FILE}.md5"
-  local CSV_URL="https://download.maxmind.com/app/geoip_download?edition_id=GeoLite2-Country-CSV&license_key=${LICENSE_KEY}&suffix=zip"
-  local MD5_URL="${CSV_URL}.md5"
+  local SHA256_FILE="${ZIPPED_FILE}.sha256"
+  local CSV_URL="https://download.maxmind.com/geoip/databases/GeoLite2-Country-CSV/download?suffix=zip"
+  local SHA256_URL="${CSV_URL}.sha256"
 
   # download files
-  curl --silent --location --output $ZIPPED_FILE "$CSV_URL" || error "Failed to download: $CSV_URL"
-  curl --silent --location --output $MD5_FILE "$MD5_URL" || error "Failed to download: $MD5_URL"
+  curl --silent --location --user "${LICENSE_KEY}" --output $ZIPPED_FILE "$CSV_URL" || error "Failed to download: $CSV_URL"
+  curl --silent --location --user "${LICENSE_KEY}" --output $SHA256_FILE "$SHA256_URL" || error "Failed to download: $SHA256_URL"
 
   # validate checksum
-  # .md5 file is not in expected format so 'md5sum --check $MD5_FILE' doesn't work
-  [[ "$(cat ${MD5_FILE})" == "$(md5sum ${ZIPPED_FILE} | awk '{print $1}')" ]] || error "Downloaded md5 checksum does not match local md5sum"
+  # .sha256 file is not in expected format so 'sha256sum --check $SHA256_FILE' doesn't work
+  [[ "$(cat ${SHA256_FILE} | awk '{print $1}')" == "$(cat ${ZIPPED_FILE} | sha256sum | awk '{print $1}')" ]] || error \
+    "Downloaded sha256 checksum does not match local sha256sum.\nCheck your License key!"
 
   # unzip into current working directory
   unzip -j -q -d . ${ZIPPED_FILE}
@@ -77,8 +78,15 @@ function build_ipv4_sets {
   readonly IPV4_IPSET_DIR="./geoipsets/ipset/ipv4/"
   readonly IPV4_NFTSET_DIR="./geoipsets/nftset/ipv4/"
 
-  rm -rf $IPV4_IPSET_DIR $IPV4_NFTSET_DIR
-  mkdir --parent $IPV4_IPSET_DIR $IPV4_NFTSET_DIR
+  if [[ $IPTABLES = "yes" ]]; then
+    rm -rf $IPV4_IPSET_DIR
+    mkdir --parent $IPV4_IPSET_DIR
+  fi
+
+  if [[ ! -v NFTABLES || $NFTABLES = "yes" ]]; then
+    rm -rf $IPV4_NFTSET_DIR
+    mkdir --parent $IPV4_NFTSET_DIR
+  fi
 
   OIFS=$IFS
   IFS=','
@@ -98,36 +106,46 @@ function build_ipv4_sets {
     SUBNET="${LINE[0]}"
     SET_NAME="${CC}.ipv4"
 
-    #
-    # iptables/ipsets
-    #
-    IPSET_FILE="${IPV4_IPSET_DIR}${SET_NAME}"
+    if [[ $IPTABLES = "yes" ]]; then
+
+      #
+      # iptables/ipsets
+      #
+
+      IPSET_FILE="${IPV4_IPSET_DIR}${SET_NAME}"
 
-    #create ipset file if it doesn't exist
-    if [[ ! -f $IPSET_FILE ]]; then
-      echo "create $SET_NAME hash:net maxelem 131072 comment" > $IPSET_FILE
+      #create ipset file if it doesn't exist
+      if [[ ! -f $IPSET_FILE ]]; then
+        echo "create $SET_NAME hash:net maxelem 131072 comment" > $IPSET_FILE
+      fi
+      echo "add ${SET_NAME} ${SUBNET} comment ${CC}" >> $IPSET_FILE
     fi
-    echo "add ${SET_NAME} ${SUBNET} comment ${CC}" >> $IPSET_FILE
 
-    #
-    # nftables set
-    #
-    NFTSET_FILE="${IPV4_NFTSET_DIR}${SET_NAME}"
+    if [[ ! -v NFTABLES || $NFTABLES = "yes" ]]; then
+
+      #
+      # nftables set
+      #
+
+      NFTSET_FILE="${IPV4_NFTSET_DIR}${SET_NAME}"
 
-    #create nft set file if it doesn't exist
-    if [[ ! -f $NFTSET_FILE ]]; then
-      echo "define $SET_NAME = {" > $NFTSET_FILE
+      #create nft set file if it doesn't exist
+      if [[ ! -f $NFTSET_FILE ]]; then
+        echo "define $SET_NAME = {" > $NFTSET_FILE
+      fi
+      echo "${SUBNET}," >> $NFTSET_FILE
     fi
-    echo "${SUBNET}," >> $NFTSET_FILE
 
   done < <(sed -e 1d "${TEMPDIR}/${ID_IPv4_RANGE_MAP}")
   IFS=$OIFS
 
   #end nft set -- better way?
-  for f in $(ls $IPV4_NFTSET_DIR)
-  do
-    echo "}" >> "${IPV4_NFTSET_DIR}$f"
-  done
+  if [[ ! -v NFTABLES || $NFTABLES = "yes" ]]; then
+    for f in "${IPV4_NFTSET_DIR}"*.ipv4
+      do
+        echo "}" >> "$f"
+      done
+  fi
 }
 
 # output
@@ -137,9 +155,16 @@ function build_ipv6_sets {
 
   readonly IPV6_IPSET_DIR="./geoipsets/ipset/ipv6/"
   readonly IPV6_NFTSET_DIR="./geoipsets/nftset/ipv6/"
+
+  if [[ $IPTABLES = "yes" ]]; then
+    rm -rf $IPV6_IPSET_DIR
+    mkdir --parent $IPV6_IPSET_DIR
+  fi
 
-  rm -rf $IPV6_IPSET_DIR $IPV6_NFTSET_DIR
-  mkdir --parent $IPV6_IPSET_DIR $IPV6_NFTSET_DIR
+  if [[ ! -v NFTABLES || $NFTABLES = "yes" ]]; then
+    rm -rf $IPV6_NFTSET_DIR
+    mkdir --parent $IPV6_NFTSET_DIR
+  fi
 
   OIFS=$IFS
   IFS=','
@@ -159,43 +184,52 @@ function build_ipv6_sets {
     SUBNET="${LINE[0]}"
     SET_NAME="${CC}.ipv6"
 
-    #
-    # iptables/ipsets
-    #
-    IPSET_FILE="${IPV6_IPSET_DIR}${SET_NAME}"
+    if [[ $IPTABLES = "yes" ]]; then
 
-    #create ipset file if it doesn't exist
-    if [[ ! -f $IPSET_FILE ]]; then
-      echo "create $SET_NAME hash:net family inet6 comment" > $IPSET_FILE
-    fi
-    echo "add ${SET_NAME} ${SUBNET} comment ${CC}" >> $IPSET_FILE
+      #
+      # iptables/ipsets
+      #
 
-    #
-    # nftables set
-    #
-    NFTSET_FILE="${IPV6_NFTSET_DIR}${SET_NAME}"
+      IPSET_FILE="${IPV6_IPSET_DIR}${SET_NAME}"
 
-    #create nft set file if it doesn't exist
-    if [[ ! -f $NFTSET_FILE ]]; then
-      echo "define $SET_NAME = {" > $NFTSET_FILE
+      #create ipset file if it doesn't exist
+      if [[ ! -f $IPSET_FILE ]]; then
+        echo "create $SET_NAME hash:net family inet6 comment" > $IPSET_FILE
+      fi
+      echo "add ${SET_NAME} ${SUBNET} comment ${CC}" >> $IPSET_FILE
     fi
-    echo "${SUBNET}," >> $NFTSET_FILE
 
+    if [[ ! -v NFTABLES || $NFTABLES = "yes" ]]; then
+
+      #
+      # nftables set
+      #
+
+      NFTSET_FILE="${IPV6_NFTSET_DIR}${SET_NAME}"
+
+      #create nft set file if it doesn't exist
+      if [[ ! -f $NFTSET_FILE ]]; then
+        echo "define $SET_NAME = {" > $NFTSET_FILE
+      fi
+      echo "${SUBNET}," >> $NFTSET_FILE
+    fi
   done < <(sed -e 1d "${TEMPDIR}/${ID_IPv6_RANGE_MAP}")
   IFS=$OIFS
 
   #end nft set -- better way?
-  for f in $(ls $IPV6_NFTSET_DIR)
-  do
-    echo "}" >> "${IPV6_NFTSET_DIR}$f"
-  done
+  if [[ ! -v NFTABLES || $NFTABLES = "yes" ]]; then
+    for f in "${IPV6_NFTSET_DIR}"*.ipv6
+      do
+        echo "}" >> "$f"
+      done
+  fi
 }
 
 # accept an optional -k switch with argument
 function main() {
 
   # get license key
-  source /etc/bcs.conf > /dev/null 2>&1
+  source ./bcs.conf > /dev/null 2>&1
   local usage="Usage: ./build-country-sets.sh [-k <LICENSE_KEY>]" 
   while getopts ":k:" opt; do
     case ${opt} in
@@ -211,8 +245,8 @@ function main() {
     esac
   done
   shift $((OPTIND -1))
-
-  [[ -z "${LICENSE_KEY}" ]] && error "No valid license key provided."
+  [[ -z "${LICENSE_KEY}" ]] && error "No license key provided.";
+   
 
   # setup
   check_progs
@@ -224,8 +258,8 @@ function main() {
   build_id_name_map
   # place set output in current working directory
   popd > /dev/null 2>&1
-  build_ipv4_sets
-  build_ipv6_sets
+  [[ ! -v IPv4 || $IPv4 = "yes" ]] && build_ipv4_sets
+  [[ ! -v IPv6 || $IPv6 = "yes" ]] && build_ipv6_sets
 }
 
 main "$@"
diff --git a/python/geoipsets.conf b/python/geoipsets.conf
@@ -1,4 +1,6 @@
 [general]
+# specify a directory where geoipsets should be saved
+output-dir=/tmp
 # list of providers from which to acquire IP ranges
 # options are:
 # 'maxmind': www.maxmind.com
@@ -28,4 +30,5 @@ address-family=ipv4,ipv6
 [maxmind]
 # specify MaxMind license key needed to download data
 # required for provider type 'maxmind', ignored by other provider types
+account-id=098765
 license-key=ABCDEFTHIJKLMNOP
diff --git a/python/geoipsets/VERSION b/python/geoipsets/VERSION
@@ -1 +1 @@
-2.3.6
+2.4.0
diff --git a/python/geoipsets/__main__.py b/python/geoipsets/__main__.py
@@ -88,6 +88,8 @@ def get_config(cli_args=None):
 
     # set defaults
     default_options = dict()
+    default_options['general'] = {''}
+    default_options['output-dir'] = default_output_dir
     default_options['provider'] = {'dbip'}
     default_options['firewall'] = {utils.Firewall.NF_TABLES.value}
     default_options['address-family'] = {utils.AddressFamily.IPV4.value}
@@ -102,34 +104,40 @@ def get_config(cli_args=None):
         if (config_file := get_config_parser(config_file_path)) is None:
             valid_conf_file = False
 
-    # step 2: provider
+    if valid_conf_file and config_file.has_section('general'):
+        general = config_file['general']
+    else:
+        valid_conf_file = False
+
+    # step 2: output_dir
+    if (output_dir := parser.parse_args(cli_args).output_dir) is not None:
+        options['output-dir'] = output_dir
+    else:
+        if valid_conf_file and (output_dir := general.get('output-dir')) is not None:
+            options['output-dir'] = output_dir
+
+    # step 3: provider
     if (providers := parser.parse_args(cli_args).provider) is not None:
         options['provider'] = set(providers)
     else:
-        if valid_conf_file and config_file.has_section('general'):
-            general = config_file['general']
-            if (providers := general.get('provider')) is not None:
-                options['provider'] = set(providers.split(','))
+        if valid_conf_file and (providers := general.get('provider')) is not None:
+            options['provider'] = set(providers.split(','))
 
-    # step 3: firewall
+    # step 4: firewall
     if (firewalls := parser.parse_args(cli_args).firewall) is not None:
         options['firewall'] = set(firewalls)
     else:
-        if valid_conf_file and config_file.has_section('general'):
-            general = config_file['general']
-            if (firewalls := general.get('firewall')) is not None:
-                options['firewall'] = set(firewalls.split(','))
+        if valid_conf_file and (firewalls := general.get('firewall')) is not None:
+            options['firewall'] = set(firewalls.split(','))
 
-    # step 4: address family
+    # step 5: address family
     if (address_family := parser.parse_args(cli_args).address_family) is not None:
         options['address-family'] = set(address_family)
     else:
-        if valid_conf_file and config_file.has_section('general'):
-            general = config_file['general']
-            if (address_family := general.get('address-family')) is not None:
-                options['address-family'] = set(address_family.split(','))
+        if valid_conf_file and (address_family := general.get('address-family')) is not None:
+            options['address-family'] = set(address_family.split(','))
 
-    # step 5: countries
+    # step 6: countries
     if (country_arg := parser.parse_args(cli_args).countries) is not None:
         country_set = set()
         try:
@@ -156,15 +164,13 @@ def get_config(cli_args=None):
             if len(countries) > 0:
                 options['countries'] = countries
 
-    # step 6: provider options
+    # step 7: provider options
     if valid_conf_file:
         for p in options.get('provider'):
             if config_file.has_section(p):
                 provider_options = config_file[p]
                 options[p] = provider_options
 
-    # step 7: output path
-    options['output-dir'] = parser.parse_args(cli_args).output_dir
     return options
 
 

diff --git a/python/geoipsets/dbip.py b/python/geoipsets/dbip.py
@@ -196,6 +196,6 @@ def check_checksum(self, csv_file_bytes):
 
         # compare downloaded sha1 hash with computed version
         if expected_sha1sum != computed_sha1sum:
-            raise RuntimeError("Computed CSV file digest '{0}' does not match expected value '{1}'".format(
+            raise SystemExit("ERROR: Computed CSV file digest '{0}' does not match expected value '{1}'".format(
                 computed_sha1sum, expected_sha1sum
             ))