# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
#
# Security patch for CVE-2026-XXXXX (path traversal vulnerability leading to RCE)
# Author: Mohammed Tanveer (threatpointer)
# Date: 2026-01-12
#
import logging
import datetime
import asyncio
import os
import uuid
import re
import pandas as pd
from typing import Dict, Any, List, Union
from iotdb.Session import Session
from iotdb.SessionPool import SessionPool, PoolConfig
from iotdb.utils.SessionDataSet import SessionDataSet
from iotdb.table_session import TableSession
from iotdb.table_session_pool import TableSessionPool, TableSessionPoolConfig
from fastmcp import FastMCP
from mcp.types import TextContent
from iotdb_mcp_server.config import Config
# Initialize FastMCP server
mcp = FastMCP("iotdb_mcp_server")
# Configure logging
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("iotdb_mcp_server")
config = Config.from_env_arguments()
db_config = {
"host": config.host,
"port": config.port,
"user": config.user,
"password": config.password,
"database": config.database,
"sql_dialect": config.sql_dialect,
"export_path": config.export_path,
}
max_pool_size = 100 # Increased from 100 for better concurrency
logger.info(f"IoTDB Config: {db_config}")
# Ensure export directory exists
if not os.path.exists(config.export_path):
try:
os.makedirs(config.export_path)
logger.info(f"Created export directory: {config.export_path}")
except Exception as e:
logger.warning(
f"Failed to create export directory {config.export_path}: {str(e)}"
)
def sanitize_filename(filename: str, base_dir: str) -> str:
"""
Sanitize and validate filename to prevent path traversal attacks.
Security patch for CVE-2026-XXXXX
Author: Mohammed Tanveer (threatpointer)
Date: 2026-01-12
Args:
filename: The user-provided filename
base_dir: The base directory for exports (must be absolute path)
Returns:
The sanitized absolute filepath
Raises:
ValueError: If the filename contains invalid characters or attempts path traversal
Security measures:
- Rejects any path separators or traversal sequences before processing
- Validates allowed characters (alphanumeric, underscore, hyphen, dot)
- Resolves absolute path and verifies it stays within base_dir boundary
- Prevents directory traversal, symlink attacks, and path manipulation
"""
if not filename:
raise ValueError("Filename cannot be empty")
# First, reject any path separators or traversal patterns before processing
# This catches attacks before os.path.basename can process them
if "/" in filename or "\\" in filename or ".." in filename:
raise ValueError(
"Invalid filename: path separators and directory traversal sequences are not allowed"
)
# Remove any directory components - only keep the base filename
# (This is now a safety check since we already blocked separators)
filename = os.path.basename(filename)
# Validate characters - only allow alphanumeric, underscore, hyphen, and dot
if not re.match(r"^[a-zA-Z0-9_\-\.]+$", filename):
raise ValueError(
"Invalid filename: only alphanumeric characters, underscore, hyphen, and dot are allowed"
)
# Prevent filenames that are just dots or empty after sanitization
if not filename or filename in (".", ".."):
raise ValueError("Invalid filename")
# Prevent filenames that start with a dot followed by a dot (like ..something)
if filename.startswith(".."):
raise ValueError("Invalid filename: cannot start with '..'")
# Construct the full path
filepath = os.path.join(base_dir, filename)
# Resolve to absolute path (resolves symlinks and normalizes path)
filepath_real = os.path.realpath(filepath)
basedir_real = os.path.realpath(base_dir)
# Ensure the resolved path is within the base directory boundary
# This prevents path traversal even with symlinks or complex path manipulation
if (
not filepath_real.startswith(basedir_real + os.sep)
and filepath_real != basedir_real
):
raise ValueError(
f"Path traversal detected: file must be within export directory"
)
return filepath_real
if config.sql_dialect == "tree":
# Configure connection pool with optimized settings
pool_config = PoolConfig(
node_urls=[str(config.host) + ":" + str(config.port)],
user_name=config.user,
password=config.password,
fetch_size=1024, # Fetch size for queries
time_zone="UTC+8", # Consistent timezone
max_retry=3, # Connection retry attempts
)
# Optimize pool size based on expected concurrent queries
wait_timeout_in_ms = 5000 # Increased from 3000 for better reliability
session_pool = SessionPool(pool_config, max_pool_size, wait_timeout_in_ms)
@mcp.tool()
async def metadata_query(query_sql: str) -> list[TextContent]:
"""Execute metadata queries on IoTDB to explore database structure and statistics.
Args:
query_sql: The metadata query to execute. Supported queries:
- SHOW DATABASES [path]: List all databases or databases under a specific path
- SHOW TIMESERIES [path]: List all time series or time series under a specific path
- SHOW CHILD PATHS [path]: List child paths under a specific path
- SHOW CHILD NODES [path]: List child nodes under a specific path
- SHOW DEVICES [path]: List all devices or devices under a specific path
- COUNT TIMESERIES [path]: Count time series under a specific path
- COUNT NODES [path]: Count nodes under a specific path
- COUNT DEVICES [path]: Count devices under a specific path
- if path is not provided, the query will be applied to root.**
Examples:
SHOW DATABASES root.**
SHOW TIMESERIES root.ln.**
SHOW CHILD PATHS root.ln
SHOW CHILD PATHS root.ln.*.*
SHOW CHILD NODES root.ln
SHOW DEVICES root.ln.**
COUNT TIMESERIES root.ln.**
COUNT NODES root.ln
COUNT DEVICES root.ln
"""
session = None
try:
session = session_pool.get_session()
stmt = query_sql.strip().upper()
# Process SHOW DATABASES
if (
stmt.startswith("SHOW DATABASES")
or stmt.startswith("SHOW TIMESERIES")
or stmt.startswith("SHOW CHILD PATHS")
or stmt.startswith("SHOW CHILD NODES")
or stmt.startswith("SHOW DEVICES")
or stmt.startswith("COUNT TIMESERIES")
or stmt.startswith("COUNT NODES")
or stmt.startswith("COUNT DEVICES")
):
res = session.execute_query_statement(query_sql)
return prepare_res(res, session)
else:
session.close()
raise ValueError(
"Unsupported metadata query. Please use one of the supported query types."
)
except Exception as e:
if session:
session.close()
logger.error(f"Failed to execute metadata query: {str(e)}")
raise
@mcp.tool()
async def select_query(query_sql: str) -> list[TextContent]:
"""Execute a SELECT query on the IoTDB tree SQL dialect.
Args:
query_sql: The SQL query to execute (using TREE dialect, time using ISO 8601 format, e.g. 2017-11-01T00:08:00.000).
SQL Syntax:
SELECT [LAST] selectExpr [, selectExpr] ...
[INTO intoItem [, intoItem] ...]
FROM prefixPath [, prefixPath] ...
[WHERE whereCondition]
[GROUP BY {
([startTime, endTime), interval [, slidingStep]) |
LEVEL = levelNum [, levelNum] ... |
TAGS(tagKey [, tagKey] ... |
VARIATION(expression[,delta][,ignoreNull=true/false]) |
CONDITION(expression,[keep>/>=/=/</<=]threshold[,ignoreNull=true/false]) |
SESSION(timeInterval) |
COUNT(expression, size[,ignoreNull=true/false])
}]
[HAVING havingCondition]
[ORDER BY sortKey {ASC | DESC}]
[FILL ({PREVIOUS | LINEAR | constant}) (, interval=DURATION_LITERAL)?)]
[SLIMIT seriesLimit] [SOFFSET seriesOffset]
[LIMIT rowLimit] [OFFSET rowOffset]
[ALIGN BY {TIME | DEVICE}]
Examples:
select temperature from root.ln.wf01.wt01 where time < 2017-11-01T00:08:00.000
select status, temperature from root.ln.wf01.wt01 where (time > 2017-11-01T00:05:00.000 and time < 2017-11-01T00:12:00.000) or (time >= 2017-11-01T16:35:00.000 and time <= 2017-11-01T16:37:00.000)
select * from root.ln.** where time > 1 order by time desc limit 10;
Supported Aggregate Functions:
SUM
COUNT
MAX_VALUE
MIN_VALUE
AVG
VARIANCE
MAX_TIME
MIN_TIME
...
"""
session = None
try:
session = session_pool.get_session()
stmt = query_sql.strip().upper()
# Regular SELECT queries
if stmt.startswith("SELECT"):
res = session.execute_query_statement(query_sql)
return prepare_res(res, session)
else:
session.close()
raise ValueError("Only SELECT queries are allowed for select_query")
except Exception as e:
if session:
session.close()
logger.error(f"Failed to execute select query: {str(e)}")
raise
@mcp.tool()
async def export_query(
query_sql: str, format: str = "csv", filename: str = None
) -> list[TextContent]:
"""Execute a query and export the results to a CSV or Excel file.
Args:
query_sql: The SQL query to execute (using TREE dialect, time using ISO 8601 format, e.g. 2017-11-01T00:08:00.000)
format: Export format, either "csv" or "excel" (default: "csv")
filename: Optional filename for the exported file. If not provided, a unique filename will be generated.
SQL Syntax:
SELECT ⟨select_list⟩
FROM ⟨tables⟩
[WHERE ⟨condition⟩]
[GROUP BY ⟨groups⟩]
[HAVING ⟨group_filter⟩]
[FILL ⟨fill_methods⟩]
[ORDER BY ⟨order_expression⟩]
[OFFSET ⟨n⟩]
[LIMIT ⟨n⟩];
Returns:
Information about the exported file and a preview of the data (max 10 rows)
"""
session = None
try:
session = session_pool.get_session()
stmt = query_sql.strip().upper()
if stmt.startswith("SELECT") or stmt.startswith("SHOW"):
# Execute the query
res = session.execute_query_statement(query_sql)
# Create a pandas DataFrame
df = res.todf()
# Close the session
session.close()
# Generate unique filename with timestamp
timestamp = int(datetime.datetime.now().timestamp())
if filename is None:
# Generate a unique filename if not provided
filename = f"dump_{uuid.uuid4().hex[:4]}_{timestamp}"
if format.lower() == "csv":
if filename.lower().endswith(".csv"):
filename = filename[:-4]
# Sanitize filename to prevent path traversal attacks
filepath = sanitize_filename(f"{filename}.csv", config.export_path)
df.to_csv(filepath, index=False)
elif format.lower() == "excel":
if filename.lower().endswith(".xlsx"):
filename = filename[:-5]
# Sanitize filename to prevent path traversal attacks
filepath = sanitize_filename(f"{filename}.xlsx", config.export_path)
df.to_excel(filepath, index=False)
else:
raise ValueError("Format must be either 'csv' or 'excel'")
# Generate preview (first 10 rows)
preview_rows = min(10, len(df))
preview_data = []
preview_data.append(",".join(df.columns)) # Header
for i in range(preview_rows):
preview_data.append(",".join(map(str, df.iloc[i])))
# Return information
return [
TextContent(
type="text",
text=f"Query results exported to {filepath}\n\nPreview (first {preview_rows} rows):\n"
+ "\n".join(preview_data),
)
]
else:
raise ValueError("Only SELECT or SHOW queries are allowed for export")
except Exception as e:
if session:
session.close()
logger.error(f"Failed to export query: {str(e)}")
raise
def prepare_res(_res: SessionDataSet, _session: Session) -> list[TextContent]:
columns = _res.get_column_names()
result = []
while _res.has_next():
record = _res.next()
if columns[0] == "Time":
timestamp = record.get_timestamp()
# formatted_time = datetime.datetime.fromtimestamp(timestamp/1000).strftime('%Y-%m-%dT%H:%M:%S.%f')[:-3]
row = record.get_fields()
result.append(str(timestamp) + "," + ",".join(map(str, row)))
else:
row = record.get_fields()
result.append(",".join(map(str, row)))
_session.close()
return [
TextContent(
type="text",
text="\n".join([",".join(columns)] + result),
)
]
elif config.sql_dialect == "table":
session_pool_config = TableSessionPoolConfig(
node_urls=[str(config.host) + ":" + str(config.port)],
username=config.user,
password=config.password,
max_pool_size=max_pool_size, # Increased from 5 for better concurrency
database=None if len(config.database) == 0 else config.database,
)
session_pool = TableSessionPool(session_pool_config)
@mcp.tool()
async def read_query(query_sql: str) -> list[TextContent]:
"""Execute a SELECT query on the IoTDB. Please use table sql_dialect when generating SQL queries.
Args:
query_sql: The SQL query to execute (using TABLE dialect, time using ISO 8601 format, e.g. 2017-11-01T00:08:00.000)
"""
table_session = None
try:
table_session = session_pool.get_session()
stmt = query_sql.strip().upper()
# Regular SELECT queries
if (
stmt.startswith("SELECT")
or stmt.startswith("DESCRIBE")
or stmt.startswith("SHOW")
):
res = table_session.execute_query_statement(query_sql)
return prepare_res(res, table_session)
else:
table_session.close()
raise ValueError("Only SELECT queries are allowed for read_query")
except Exception as e:
if table_session:
table_session.close()
logger.error(f"Failed to execute query: {str(e)}")
raise
@mcp.tool()
async def list_tables() -> list[TextContent]:
"""List all tables in the IoTDB database."""
table_session = None
try:
table_session = session_pool.get_session()
res = table_session.execute_query_statement("SHOW TABLES")
result = ["Tables_in_" + db_config["database"]] # Header
while res.has_next():
result.append(str(res.next().get_fields()[0]))
table_session.close()
return [TextContent(type="text", text="\n".join(result))]
except Exception as e:
if table_session:
table_session.close()
logger.error(f"Failed to list tables: {str(e)}")
raise
@mcp.tool()
async def describe_table(table_name: str) -> list[TextContent]:
"""Get the schema information for a specific table
Args:
table_name: name of the table to describe
"""
table_session = None
try:
table_session = session_pool.get_session()
res = table_session.execute_query_statement(
"DESC " + table_name + " details"
)
return prepare_res(res, table_session)
except Exception as e:
if table_session:
table_session.close()
logger.error(f"Failed to describe table {table_name}: {str(e)}")
raise
@mcp.tool()
async def export_table_query(
query_sql: str, format: str = "csv", filename: str = None
) -> list[TextContent]:
"""Execute a query and export the results to a CSV or Excel file.
Args:
query_sql: The SQL query to execute (using TABLE dialect, time using ISO 8601 format, e.g. 2017-11-01T00:08:00.000)
format: Export format, either "csv" or "excel" (default: "csv")
filename: Optional filename for the exported file. If not provided, a unique filename will be generated.
SQL Syntax:
SELECT ⟨select_list⟩
FROM ⟨tables⟩
[WHERE ⟨condition⟩]
[GROUP BY ⟨groups⟩]
[HAVING ⟨group_filter⟩]
[FILL ⟨fill_methods⟩]
[ORDER BY ⟨order_expression⟩]
[OFFSET ⟨n⟩]
[LIMIT ⟨n⟩];
Returns:
Information about the exported file and a preview of the data (max 10 rows)
"""
table_session = None
try:
table_session = session_pool.get_session()
stmt = query_sql.strip().upper()
if (
stmt.startswith("SELECT")
or stmt.startswith("SHOW")
or stmt.startswith("DESCRIBE")
or stmt.startswith("DESC")
):
# Execute the query
res = table_session.execute_query_statement(query_sql)
# Create a pandas DataFrame
df = res.todf()
# Close the session
table_session.close()
# Generate unique filename with timestamp
timestamp = int(datetime.datetime.now().timestamp())
if filename is None:
filename = f"dump_{uuid.uuid4().hex[:4]}_{timestamp}"
if format.lower() == "csv":
if filename.lower().endswith(".csv"):
filename = filename[:-4]
# Sanitize filename to prevent path traversal attacks
filepath = sanitize_filename(f"{filename}.csv", config.export_path)
df.to_csv(filepath, index=False)
elif format.lower() == "excel":
if filename.lower().endswith(".xlsx"):
filename = filename[:-5]
# Sanitize filename to prevent path traversal attacks
filepath = sanitize_filename(f"{filename}.xlsx", config.export_path)
df.to_excel(filepath, index=False)
else:
raise ValueError("Format must be either 'csv' or 'excel'")
# Generate preview (first 10 rows)
preview_rows = min(10, len(df))
preview_data = []
preview_data.append(",".join(df.columns)) # Header
for i in range(preview_rows):
preview_data.append(",".join(map(str, df.iloc[i])))
# Return information
return [
TextContent(
type="text",
text=f"Query results exported to {filepath}\n\nPreview (first {preview_rows} rows):\n"
+ "\n".join(preview_data),
)
]
else:
raise ValueError(
"Only SELECT, SHOW or DESCRIBE queries are allowed for export"
)
except Exception as e:
if table_session:
table_session.close()
logger.error(f"Failed to export table query: {str(e)}")
raise
def prepare_res(
_res: SessionDataSet, _table_session: TableSession
) -> list[TextContent]:
columns = _res.get_column_names()
result = []
while _res.has_next():
row = _res.next().get_fields()
result.append(",".join(map(str, row)))
_table_session.close()
return [
TextContent(
type="text",
text="\n".join([",".join(columns)] + result),
)
]
def main():
logger.info("iotdb_mcp_server running with stdio transport")
# Initialize and run the server
mcp.run(transport="stdio")
if __name__ == "__main__":
main()
