Coverage for src/lexigram/graphql/security/alias.py: 100%

1"""Query alias limiting. 

2 

3This module provides alias limiting for GraphQL queries 

4to prevent abuse through excessive field aliasing. 

5""" 

6 

7from __future__ import annotations 

8 

9from typing import TYPE_CHECKING 

10 

11from strawberry.extensions import SchemaExtension 

12 

13from lexigram.graphql.exceptions import GraphQLError 

14from lexigram.logging import get_logger 

15 

16if TYPE_CHECKING: 

17 from collections.abc import Iterator 

18 

19 from graphql.language.ast import DocumentNode, SelectionSetNode 

20 

21 

22logger = get_logger(__name__) 

23 

24 

25class AliasLimitValidator: 

26 """Validate query alias count against a limit. 

27 

28 Analyzes GraphQL queries to count aliases and validates 

29 against a configured maximum. 

30 

31 Example: 

32 ```python 

33 validator = AliasLimitValidator(max_aliases=10) 

34 

35 count = validator.count_aliases(document) 

36 validator.validate(document) # Raises if too many 

37 ``` 

38 """ 

39 

40 def __init__(self, max_aliases: int = 10) -> None: 

41 """Initialize the validator. 

42 

43 Args: 

44 max_aliases: Maximum allowed aliases. 

45 """ 

46 self._max_aliases = max_aliases 

47 

48 @property 

49 def max_aliases(self) -> int: 

50 """Get maximum alias limit.""" 

51 return self._max_aliases 

52 

53 def count_aliases(self, document: DocumentNode) -> int: 

54 """Count the number of aliases in a query. 

55 

56 Args: 

57 document: Parsed GraphQL document. 

58 

59 Returns: 

60 Number of aliases. 

61 """ 

62 count = 0 

63 

64 for definition in document.definitions: 

65 if hasattr(definition, "selection_set"): 

66 count += self._count_in_selection_set(definition.selection_set) 

67 

68 return count 

69 

70 def _count_in_selection_set( 

71 self, 

72 selection_set: SelectionSetNode | None, 

73 ) -> int: 

74 """Count aliases in a selection set. 

75 

76 Args: 

77 selection_set: Selection set to analyze. 

78 

79 Returns: 

80 Number of aliases. 

81 """ 

82 if selection_set is None: 

83 return 0 

84 

85 count = 0 

86 

87 for selection in selection_set.selections: 

88 # Check for alias 

89 if hasattr(selection, "alias") and selection.alias: 

90 count += 1 

91 

92 # Recurse into nested selections 

93 if hasattr(selection, "selection_set"): 

94 count += self._count_in_selection_set(selection.selection_set) 

95 

96 return count 

97 

98 def validate(self, document: DocumentNode) -> None: 

99 """Validate alias count. 

100 

101 Args: 

102 document: Parsed GraphQL document. 

103 

104 Raises: 

105 SecurityError: If query has too many aliases. 

106 """ 

107 count = self.count_aliases(document) 

108 

109 if count > self._max_aliases: 

110 error = GraphQLError( 

111 f"Query has {count} aliases, exceeding maximum of {self._max_aliases}", 

112 ) 

113 error.safe = True 

114 raise error 

115 

116 logger.debug("Query has %d aliases (limit: %d)", count, self._max_aliases) 

117 

118 

119class AliasLimitExtension(SchemaExtension): 

120 """Strawberry extension for query alias limiting. 

121 

122 Add this extension to your schema to automatically 

123 validate alias count before execution. 

124 

125 Example: 

126 ```python 

127 from lexigram.graphql.security import AliasLimitExtension 

128 

129 schema = strawberry.Schema( 

130 query=Query, 

131 extensions=[AliasLimitExtension(max_aliases=10)], 

132 ) 

133 ``` 

134 """ 

135 

136 def __init__(self, max_aliases: int = 10) -> None: 

137 """Initialize the extension. 

138 

139 Args: 

140 max_aliases: Maximum allowed aliases. 

141 """ 

142 self._validator = AliasLimitValidator(max_aliases=max_aliases) 

143 

144 def on_validate(self) -> Iterator[None]: 

145 """Hook called during operation validation.""" 

146 execution_context = self.execution_context 

147 

148 # Validate alias count 

149 if execution_context.graphql_document: 

150 self._validator.validate(execution_context.graphql_document) 

151 

152 yield 

153 

154 

155__all__ = ["AliasLimitExtension", "AliasLimitValidator"]