diff --git a/airflow/providers/apache/cassandra/hooks/cassandra.py b/airflow/providers/apache/cassandra/hooks/cassandra.py index 999782e0da6a1..66d3d13aa8c25 100644 --- a/airflow/providers/apache/cassandra/hooks/cassandra.py +++ b/airflow/providers/apache/cassandra/hooks/cassandra.py @@ -18,6 +18,7 @@ """This module contains hook to integrate with Apache Cassandra.""" from __future__ import annotations +import re from typing import Any, Union from cassandra.auth import PlainTextAuthProvider @@ -188,6 +189,13 @@ def table_exists(self, table: str) -> bool: cluster_metadata = self.get_conn().cluster.metadata return keyspace in cluster_metadata.keyspaces and table in cluster_metadata.keyspaces[keyspace].tables + @staticmethod + def _sanitize_input(input_string: str) -> str: + if re.match(r"^\w+$", input_string): + return input_string + else: + raise ValueError(f"Invalid input: {input_string}") + def record_exists(self, table: str, keys: dict[str, str]) -> bool: """ Check if a record exists in Cassandra. @@ -196,9 +204,11 @@ def record_exists(self, table: str, keys: dict[str, str]) -> bool: Use dot notation to target a specific keyspace. :param keys: The keys and their values to check the existence. """ - keyspace = self.keyspace + keyspace = self._sanitize_input(self.keyspace) if self.keyspace else self.keyspace if "." in table: - keyspace, table = table.split(".", 1) + keyspace, table = map(self._sanitize_input, table.split(".", 1)) + else: + table = self._sanitize_input(table) ks_str = " AND ".join(f"{key}=%({key})s" for key in keys) query = f"SELECT * FROM {keyspace}.{table} WHERE {ks_str}" try: diff --git a/tests/integration/providers/apache/cassandra/hooks/test_cassandra.py b/tests/integration/providers/apache/cassandra/hooks/test_cassandra.py index c100e55bdb60e..c27b72d3691b2 100644 --- a/tests/integration/providers/apache/cassandra/hooks/test_cassandra.py +++ b/tests/integration/providers/apache/cassandra/hooks/test_cassandra.py @@ -17,6 +17,7 @@ # under the License. from __future__ import annotations +import re from unittest import mock import pytest @@ -232,3 +233,19 @@ def test_table_exists_with_keyspace_from_session(self): session.shutdown() hook.shutdown_cluster() + + def test_possible_sql_injection(self): + hook = CassandraHook("cassandra_default_with_schema") + session = hook.get_conn() + cqls = [ + "DROP TABLE IF EXISTS t", + "CREATE TABLE t (pk1 text, pk2 text, c text, PRIMARY KEY (pk1, pk2))", + "INSERT INTO t (pk1, pk2, c) VALUES ('foo', 'bar', 'baz')", + ] + for cql in cqls: + session.execute(cql) + + assert hook.record_exists("t", {"pk1": "foo", "pk2": "bar"}) + assert not hook.record_exists("tt", {"pk1": "foo", "pk2": "bar"}) + with pytest.raises(ValueError, match=re.escape("Invalid input: t; DROP TABLE t; SELECT * FROM t")): + hook.record_exists("t; DROP TABLE t; SELECT * FROM t", {"pk1": "foo", "pk2": "baz"})