fix: parse ABACUS binary options safely

dpgen/data/gen.py

@@ -19,6 +19,7 @@
 from dpgen import ROOT_PATH, dlog
 from dpgen.dispatcher.Dispatcher import make_submission
 from dpgen.generator.lib.abacus_scf import (
+    _parse_abacus_binary,
     get_abacus_input_parameters,
     get_abacus_STRU,
     make_abacus_scf_kpt,
@@ -578,8 +579,9 @@ def make_abacus_relax(jdata, mdata):
     )  # a dictionary in which all of the values are strings
     if "kspacing" not in standard_incar:
         if "gamma_only" in standard_incar:
-            if isinstance(standard_incar["gamma_only"], str):
-                standard_incar["gamma_only"] = int(eval(standard_incar["gamma_only"]))
+            standard_incar["gamma_only"] = _parse_abacus_binary(
+                standard_incar["gamma_only"], "gamma_only"
+            )
             if standard_incar["gamma_only"] == 0:
                 if "relax_kpt" not in jdata:
                     raise RuntimeError("Cannot find any k-points information.")
@@ -922,8 +924,9 @@ def make_abacus_md(jdata, mdata):
     #        "Cannot find any k-points information."
     if "kspacing" not in standard_incar:
         if "gamma_only" in standard_incar:
-            if isinstance(standard_incar["gamma_only"], str):
-                standard_incar["gamma_only"] = int(eval(standard_incar["gamma_only"]))
+            standard_incar["gamma_only"] = _parse_abacus_binary(
+                standard_incar["gamma_only"], "gamma_only"
+            )
             if standard_incar["gamma_only"] == 0:
                 if "md_kpt" not in jdata:
                     raise RuntimeError("Cannot find any k-points information.")

dpgen/generator/lib/abacus_scf.py

@@ -11,6 +11,22 @@
 bohr2ang = 0.52917721067
 
 
+def _parse_abacus_binary(value, field_name):
+    """Parse an ABACUS binary option without evaluating Python expressions."""
+    if isinstance(value, str):
+        normalized = value.strip().lower()
+        if normalized in {"1", "true", "t", "yes", "y", "on"}:
+            return 1
+        if normalized in {"0", "false", "f", "no", "n", "off"}:
+            return 0
+    elif isinstance(value, (bool, int, float, np.integer, np.floating)) and value in {
+        0,
+        1,
+    }:
+        return int(value)
+    raise ValueError(f"{field_name!r} should be either 0 or 1")
+
+
 def make_abacus_scf_kpt(fp_params):
     # Make KPT file for abacus pw scf calculation.
     # KPT file is the file containing k points infomation in ABACUS scf calculation.
@@ -76,10 +92,8 @@ def make_abacus_scf_input(fp_params, extra_file_path=""):
         elif key == "dft_functional":
             ret += "dft_functional {}\n".format(fp_params["dft_functional"])
         elif key == "gamma_only":
-            if isinstance(fp_params["gamma_only"], str):
-                fp_params["gamma_only"] = int(eval(fp_params["gamma_only"]))
-            assert fp_params["gamma_only"] == 0 or fp_params["gamma_only"] == 1, (
-                "'gamma_only' should be either 0 or 1."
+            fp_params["gamma_only"] = _parse_abacus_binary(
+                fp_params["gamma_only"], "gamma_only"
             )
             ret += "gamma_only %d\n" % fp_params["gamma_only"]  # noqa: UP031
         elif key == "mixing_type":
@@ -98,10 +112,8 @@ def make_abacus_scf_input(fp_params, extra_file_path=""):
             )
             ret += "mixing_beta {:f}\n".format(fp_params["mixing_beta"])
         elif key == "symmetry":
-            if isinstance(fp_params["symmetry"], str):
-                fp_params["symmetry"] = int(eval(fp_params["symmetry"]))
-            assert fp_params["symmetry"] == 0 or fp_params["symmetry"] == 1, (
-                "'symmetry' should be either 0 or 1."
+            fp_params["symmetry"] = _parse_abacus_binary(
+                fp_params["symmetry"], "symmetry"
             )
             ret += "symmetry %d\n" % fp_params["symmetry"]  # noqa: UP031
         elif key == "nbands":
@@ -150,29 +162,20 @@ def make_abacus_scf_input(fp_params, extra_file_path=""):
             )
             ret += "smearing_sigma {:f}\n".format(fp_params["smearing_sigma"])
         elif key == "cal_force":
-            if isinstance(fp_params["cal_force"], str):
-                fp_params["cal_force"] = int(eval(fp_params["cal_force"]))
-            assert fp_params["cal_force"] == 0 or fp_params["cal_force"] == 1, (
-                "'cal_force' should be either 0 or 1."
+            fp_params["cal_force"] = _parse_abacus_binary(
+                fp_params["cal_force"], "cal_force"
             )
             ret += "cal_force %d\n" % fp_params["cal_force"]  # noqa: UP031
         elif key == "cal_stress":
-            if isinstance(fp_params["cal_stress"], str):
-                fp_params["cal_stress"] = int(eval(fp_params["cal_stress"]))
-            assert fp_params["cal_stress"] == 0 or fp_params["cal_stress"] == 1, (
-                "'cal_stress' should be either 0 or 1."
+            fp_params["cal_stress"] = _parse_abacus_binary(
+                fp_params["cal_stress"], "cal_stress"
             )
             ret += "cal_stress %d\n" % fp_params["cal_stress"]  # noqa: UP031
         # paras for deepks
         elif key == "deepks_out_labels":
-            if isinstance(fp_params["deepks_out_labels"], str):
-                fp_params["deepks_out_labels"] = int(
-                    eval(fp_params["deepks_out_labels"])
-                )
-            assert (
-                fp_params["deepks_out_labels"] == 0
-                or fp_params["deepks_out_labels"] == 1
-            ), "'deepks_out_labels' should be either 0 or 1."
+            fp_params["deepks_out_labels"] = _parse_abacus_binary(
+                fp_params["deepks_out_labels"], "deepks_out_labels"
+            )
             ret += "deepks_out_labels %d\n" % fp_params["deepks_out_labels"]  # noqa: UP031
         elif key == "deepks_descriptor_lmax":
             fp_params["deepks_descriptor_lmax"] = int(
@@ -183,10 +186,8 @@ def make_abacus_scf_input(fp_params, extra_file_path=""):
             )
             ret += "deepks_descriptor_lmax %d\n" % fp_params["deepks_descriptor_lmax"]  # noqa: UP031
         elif key == "deepks_scf":
-            if isinstance(fp_params["deepks_scf"], str):
-                fp_params["deepks_scf"] = int(eval(fp_params["deepks_scf"]))
-            assert fp_params["deepks_scf"] == 0 or fp_params["deepks_scf"] == 1, (
-                "'deepks_scf' should be either 0 or 1."
+            fp_params["deepks_scf"] = _parse_abacus_binary(
+                fp_params["deepks_scf"], "deepks_scf"
             )
             ret += "deepks_scf %d\n" % fp_params["deepks_scf"]  # noqa: UP031
         elif key == "deepks_model":

tests/generator/test_abacus_safe_values.py

@@ -0,0 +1,39 @@
+import tempfile
+import unittest
+from pathlib import Path
+
+from dpgen.generator.lib.abacus_scf import (
+    _parse_abacus_binary,
+    make_abacus_scf_input,
+)
+
+
+class TestAbacusSafeValues(unittest.TestCase):
+    def test_binary_string_values(self):
+        cases = {
+            "1": 1,
+            "true": 1,
+            "T": 1,
+            "0": 0,
+            "false": 0,
+            "F": 0,
+        }
+        for value, expected in cases.items():
+            with self.subTest(value=value):
+                self.assertEqual(_parse_abacus_binary(value, "option"), expected)
+
+    def test_python_expression_is_rejected_without_execution(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            marker = Path(tmpdir) / "eval-ran"
+            payload = (
+                f"__import__('pathlib').Path({str(marker)!r}).write_text('ran') or 1"
+            )
+
+            with self.assertRaisesRegex(ValueError, "gamma_only"):
+                make_abacus_scf_input({"gamma_only": payload})
+
+            self.assertFalse(marker.exists())
+
+
+if __name__ == "__main__":
+    unittest.main()