1  # Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
 2  #
 3  # Licensed under the Apache License, Version 2.0 (the "License");
 4  # you may not use this file except in compliance with the License.
 5  # You may obtain a copy of the License at
 6  #
 7  #     http://www.apache.org/licenses/LICENSE-2.0
 8  #
 9  # Unless required by applicable law or agreed to in writing, software
10  # distributed under the License is distributed on an "AS IS" BASIS,
11  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12  # See the License for the specific language governing permissions and
13  # limitations under the License.
14  #
15  # Requirement: Any integration or derivative work must explicitly attribute
16  # Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
17  # documentation or user interface, as detailed in the NOTICE file.
18  
19  from enum import Enum
20  from typing import List, Type
21  
22  
23  def validate_vulnerability_types(
24      vulnerability_name: str, types: List[str], allowed_type: Type[Enum]
25  ) -> List[Enum]:
26      if not isinstance(types, list):
27          raise TypeError(
28              f"The 'types' attribute for the {vulnerability_name} vulnerability must be a list of strings."
29          )
30      if not types:
31          raise ValueError(
32              f"The 'types' attribute for the {vulnerability_name} vulnerability attribute cannot be an empty list."
33          )
34  
35      duplicate_types = [t for t in set(types) if types.count(t) > 1]
36      if duplicate_types:
37          quoted_duplicate_types = [f'"{t}"' for t in duplicate_types]
38          raise ValueError(
39              f"Duplicate types detected: {', '.join(quoted_duplicate_types)} for the {vulnerability_name} vulnerability."
40          )
41  
42      valid_values = [type.value for type in allowed_type]
43      invalid_types = [t for t in types if t not in valid_values]
44      if invalid_types:
45          if len(invalid_types) == 1:
46              invalid_types_text = f'"{invalid_types[0]}"'
47          elif len(invalid_types) == 2:
48              invalid_types_text = (
49                  f'"{invalid_types[0]}" and "{invalid_types[1]}"'
50              )
51          else:
52              quoted_invalid_types = [f'"{t}"' for t in invalid_types[:-1]]
53              invalid_types_text = (
54                  f'{", ".join(quoted_invalid_types)}, and "{invalid_types[-1]}"'
55              )
56  
57          if len(valid_values) == 1:
58              valid_values_text = f'"{valid_values[0]}"'
59          elif len(valid_values) == 2:
60              valid_values_text = f'"{valid_values[0]}" and "{valid_values[1]}"'
61          else:
62              quoted_valid_values = [f'"{v}"' for v in valid_values[:-1]]
63              valid_values_text = (
64                  f'{", ".join(quoted_valid_values)}, and "{valid_values[-1]}"'
65              )
66  
67          verb = "does" if len(invalid_types) == 1 else "do"
68          raise ValueError(
69              f"The {invalid_types_text} {verb} not exist for the {vulnerability_name} vulnerability. "
70              f"The available types are {valid_values_text}."
71          )
72  
73      enum_types = [
74          enum_type for enum_type in allowed_type if enum_type.value in types
75      ]
76      return enum_types