From 9c9298e3b22909a2aa3926a3b80b8aac68cacd59 Mon Sep 17 00:00:00 2001
From: Yan Cheng <58191769+yanchengnv@users.noreply.github.com>
Date: Thu, 13 Jun 2024 10:42:33 -0400
Subject: [PATCH] [2.5] Consolidate text file name validation (#2632)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* test pr

* fix formatting issues

---------

Co-authored-by: Yuan-Ting Hsieh (謝沅廷) <yuantingh@nvidia.com>
---
 nvflare/fuel/hci/client/fl_admin_api.py | 11 ++++------
 nvflare/fuel/hci/cmd_arg_utils.py       | 28 +++++++++++++++++++------
 nvflare/private/fed/server/shell_cmd.py | 13 ++++--------
 3 files changed, 30 insertions(+), 22 deletions(-)

diff --git a/nvflare/fuel/hci/client/fl_admin_api.py b/nvflare/fuel/hci/client/fl_admin_api.py
index 6ff6a63477..59fc9fef8b 100644
--- a/nvflare/fuel/hci/client/fl_admin_api.py
+++ b/nvflare/fuel/hci/client/fl_admin_api.py
@@ -25,7 +25,7 @@
 from nvflare.fuel.hci.client.api_status import APIStatus
 from nvflare.fuel.hci.client.fl_admin_api_constants import FLDetailKey
 from nvflare.fuel.hci.client.fl_admin_api_spec import APISyntaxError, FLAdminAPIResponse, FLAdminAPISpec, TargetType
-from nvflare.fuel.hci.cmd_arg_utils import get_file_extension
+from nvflare.fuel.hci.cmd_arg_utils import validate_text_file_name
 from nvflare.security.logging import secure_format_exception
 
 from .overseer_service_finder import ServiceFinderByOverseer
@@ -211,12 +211,9 @@ def _validate_file_string(self, file: str) -> str:
             if p == "..":
                 raise APISyntaxError(".. in file path is not allowed")
 
-        file_extension = get_file_extension(file)
-        if file_extension not in [".txt", ".log", ".json", ".csv", ".sh", ".config", ".py"]:
-            raise APISyntaxError(
-                "this command cannot be applied to file {}. Only files with the following extensions are "
-                "permitted: .txt, .log, .json, .csv, .sh, .config, .py".format(file)
-            )
+        err = validate_text_file_name(file)
+        if err:
+            raise APISyntaxError(err)
         return file
 
     def _validate_sp_string(self, sp_string) -> str:
diff --git a/nvflare/fuel/hci/cmd_arg_utils.py b/nvflare/fuel/hci/cmd_arg_utils.py
index 3d7f518777..2deebe6c8b 100644
--- a/nvflare/fuel/hci/cmd_arg_utils.py
+++ b/nvflare/fuel/hci/cmd_arg_utils.py
@@ -140,15 +140,31 @@ def get_file_extension(file: str) -> str:
     return ex
 
 
+def validate_text_file_name(file_name: str) -> str:
+    """Check the specified file name whether it is acceptable.
+
+    Args:
+        file_name: file name to be checked.
+
+    Returns: error string if invalid; or empty string if valid
+
+    """
+    file_extension = get_file_extension(file_name)
+    if file_extension not in [".txt", ".log", ".json", ".csv", ".sh", ".config", ".py"]:
+        return (
+            f"this command cannot be applied to file {file_name}. Only files with the following extensions are "
+            "permitted: .txt, .log, .json, .csv, .sh, .config, .py"
+        )
+    else:
+        return ""
+
+
 def validate_file_string(file: str) -> str:
     """Returns the file string if it is valid."""
     validate_path_string(file)
-    file_extension = get_file_extension(file)
-    if file_extension not in [".txt", ".log", ".json", ".csv", ".sh", ".config", ".py"]:
-        raise SyntaxError(
-            "this command cannot be applied to file {}. Only files with the following extensions are "
-            "permitted: .txt, .log, .json, .csv, .sh, .config, .py".format(file)
-        )
+    err = validate_text_file_name(file)
+    if err:
+        raise SyntaxError(err)
     return file
 
 
diff --git a/nvflare/private/fed/server/shell_cmd.py b/nvflare/private/fed/server/shell_cmd.py
index 0b82a4f889..c7b24d2f35 100644
--- a/nvflare/private/fed/server/shell_cmd.py
+++ b/nvflare/private/fed/server/shell_cmd.py
@@ -16,7 +16,7 @@
 import subprocess
 from typing import List
 
-from nvflare.fuel.hci.cmd_arg_utils import get_file_extension, join_args
+from nvflare.fuel.hci.cmd_arg_utils import join_args, validate_text_file_name
 from nvflare.fuel.hci.conn import Connection
 from nvflare.fuel.hci.proto import MetaStatusValue, make_meta
 from nvflare.fuel.hci.reg import CommandModule, CommandModuleSpec, CommandSpec
@@ -185,14 +185,9 @@ def validate_shell_command(self, args: List[str], parse_result):
                         return ".. in path name is not allowed"
 
                 if self.text_file_only:
-                    # check whether the file name is ended with numbers. If so, the actual file extension is before it.
-                    # this is the case that when the log file (log.txt) is rotated, the previous file becomes log.txt.1.
-                    file_extension = get_file_extension(f)
-                    if file_extension not in [".txt", ".log", ".json", ".csv", ".sh", ".config", ".py"]:
-                        return (
-                            "this command cannot be applied to file {}. Only files with the following extensions "
-                            "are permitted: .txt, .log, .json, .csv, .sh, .config, .py".format(f)
-                        )
+                    err = validate_text_file_name(f)
+                    if err:
+                        return err
 
         return ""