Merge pull request #413 from J08nY/feat/proxy-artifacts

Allow for download through our caching proxy at sec-certs.org

4 files changed, 38 insertions, 8 deletions

diff --git a/src/sec_certs/configuration.py b/src/sec_certs/configuration.py
index 906960c4..47ac5821 100644
--- a/src/sec_certs/configuration.py
+++ b/src/sec_certs/configuration.py
@@ -106,6 +106,8 @@ class Configuration(BaseSettings):
         ge=0,
         le=100,
     )
+    cc_use_proxy: bool = Field(False, description="Download CC artifacts through the sec-certs.org proxy.")
+    fips_use_proxy: bool = Field(False, description="Download FIPS artifacts through the sec-certs.org proxy.")
     enable_progress_bars: bool = Field(
         True, description="If true, progress bars will be printed to stdout during computation."
     )
diff --git a/src/sec_certs/sample/cc.py b/src/sec_certs/sample/cc.py
index e753eb83..9f8ba2bd 100644
--- a/src/sec_certs/sample/cc.py
+++ b/src/sec_certs/sample/cc.py
@@ -18,6 +18,7 @@ import sec_certs.utils.pdf
 import sec_certs.utils.sanitization
 from sec_certs import constants
 from sec_certs.cert_rules import SARS_IMPLIED_FROM_EAL, cc_rules, rules, security_level_csv_scan
+from sec_certs.configuration import config
 from sec_certs.sample.cc_certificate_id import canonicalize, schemes
 from sec_certs.sample.certificate import Certificate, References, logger
 from sec_certs.sample.certificate import Heuristics as BaseHeuristics
@@ -42,8 +43,7 @@ class CCCertificate(
     the certificate can handle itself. `CCDataset` class then instrument this functionality.
     """
 
-    cc_url = "http://www.commoncriteriaportal.org"
-    empty_st_url = "http://www.commoncriteriaportal.org/files/epfiles/"
+    cc_url = "https://www.commoncriteriaportal.org"
 
     @dataclass(eq=True, frozen=True)
     class MaintenanceReport(ComplexSerializableType):
@@ -753,7 +753,7 @@ class CCCertificate(
         if not cert.report_link:
             exit_code = "No link"
         else:
-            exit_code = helpers.download_file(cert.report_link, cert.state.report.pdf_path)
+            exit_code = helpers.download_file(cert.report_link, cert.state.report.pdf_path, proxy=config.cc_use_proxy)
         if exit_code != requests.codes.ok:
             error_msg = f"failed to download report from {cert.report_link}, code: {exit_code}"
             logger.error(f"Cert dgst: {cert.dgst} " + error_msg)
@@ -773,7 +773,9 @@ class CCCertificate(
         :return CCCertificate: returns the modified certificate with updated state
         """
         exit_code: str | int = (
-            helpers.download_file(cert.st_link, cert.state.st.pdf_path) if cert.st_link else "No link"
+            helpers.download_file(cert.st_link, cert.state.st.pdf_path, proxy=config.cc_use_proxy)
+            if cert.st_link
+            else "No link"
         )
 
         if exit_code != requests.codes.ok:
@@ -795,7 +797,9 @@ class CCCertificate(
         :return CCCertificate: returns the modified certificate with updated state
         """
         exit_code: str | int = (
-            helpers.download_file(cert.cert_link, cert.state.cert.pdf_path) if cert.cert_link else "No link"
+            helpers.download_file(cert.cert_link, cert.state.cert.pdf_path, proxy=config.cc_use_proxy)
+            if cert.cert_link
+            else "No link"
         )
 
         if exit_code != requests.codes.ok:
diff --git a/src/sec_certs/sample/fips.py b/src/sec_certs/sample/fips.py
index 26b9833e..cee8cb9d 100644
--- a/src/sec_certs/sample/fips.py
+++ b/src/sec_certs/sample/fips.py
@@ -558,7 +558,11 @@ class FIPSCertificate(
 
     @staticmethod
     def download_module(cert: FIPSCertificate) -> FIPSCertificate:
-        if (exit_code := helpers.download_file(cert.module_html_url, cert.state.module_html_path)) != requests.codes.ok:
+        if (
+            exit_code := helpers.download_file(
+                cert.module_html_url, cert.state.module_html_path, proxy=config.fips_use_proxy
+            )
+        ) != requests.codes.ok:
             error_msg = f"failed to download html module from {cert.module_html_url}, code {exit_code}"
             logger.error(f"Cert dgst: {cert.dgst} " + error_msg)
             cert.state.module_download_ok = False
@@ -568,7 +572,11 @@ class FIPSCertificate(
 
     @staticmethod
     def download_policy(cert: FIPSCertificate) -> FIPSCertificate:
-        if (exit_code := helpers.download_file(cert.policy_pdf_url, cert.state.policy_pdf_path)) != requests.codes.ok:
+        if (
+            exit_code := helpers.download_file(
+                cert.policy_pdf_url, cert.state.policy_pdf_path, proxy=config.fips_use_proxy
+            )
+        ) != requests.codes.ok:
             error_msg = f"failed to download pdf policy from {cert.policy_pdf_url}, code {exit_code}"
             logger.error(f"Cert dgst: {cert.dgst} " + error_msg)
             cert.state.policy_download_ok = False
diff --git a/src/sec_certs/utils/helpers.py b/src/sec_certs/utils/helpers.py
index c0440008..dedee959 100644
--- a/src/sec_certs/utils/helpers.py
+++ b/src/sec_certs/utils/helpers.py
@@ -22,11 +22,27 @@ from sec_certs.utils.tqdm import tqdm
 logger = logging.getLogger(__name__)
 
 
+_PROXIES = {
+    "https://www.commoncriteriaportal.org/": "https://sec-certs.org/proxy/cc/",
+    "https://csrc.nist.gov/": "https://sec-certs.org/proxy/fips/",
+}
+
+
 def download_file(
-    url: str, output: Path, delay: float = 0, show_progress_bar: bool = False, progress_bar_desc: str | None = None
+    url: str,
+    output: Path,
+    delay: float = 0,
+    show_progress_bar: bool = False,
+    progress_bar_desc: str | None = None,
+    proxy: bool = False,
 ) -> str | int:
     try:
         time.sleep(delay)
+        if proxy:
+            for upstream in _PROXIES:
+                if upstream in url:
+                    url = url.replace(upstream, _PROXIES[upstream])
+                    break
         # See https://github.com/psf/requests/issues/3953 for header justification
         r = requests.get(
             url,