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

1"""Query whitelist for production GraphQL endpoints. 

2 

3This module provides query whitelisting to restrict GraphQL 

4operations to a predefined set for security. 

5""" 

6 

7from __future__ import annotations 

8 

9from dataclasses import dataclass 

10 

11from lexigram import hashing # type: ignore[attr-defined] 

12 

13 

14@dataclass 

15class WhitelistEntry: 

16 """A whitelisted query entry. 

17 

18 Attributes: 

19 hash: SHA256 hash of the query. 

20 query: The actual query string. 

21 operation_name: Optional operation name. 

22 description: Human-readable description. 

23 """ 

24 

25 hash: str 

26 query: str 

27 operation_name: str | None = None 

28 description: str | None = None 

29 

30 

31class QueryWhitelist: 

32 """Query whitelist manager. 

33 

34 Example: 

35 whitelist = QueryWhitelist() 

36 whitelist.add("{ user { id name } }", "Get user") 

37 

38 # Validate incoming queries 

39 if not whitelist.is_allowed(query): 

40 raise PermissionError("Query not whitelisted") 

41 """ 

42 

43 def __init__(self, enabled: bool = False): 

44 """Initialize the whitelist. 

45 

46 Args: 

47 enabled: Whether whitelist enforcement is enabled. 

48 """ 

49 self._enabled = enabled 

50 self._entries: dict[str, WhitelistEntry] = {} 

51 

52 @property 

53 def enabled(self) -> bool: 

54 """Check if whitelist is enabled.""" 

55 return self._enabled 

56 

57 def add( 

58 self, 

59 query: str, 

60 description: str | None = None, 

61 operation_name: str | None = None, 

62 ) -> WhitelistEntry: 

63 """Add a query to the whitelist. 

64 

65 Args: 

66 query: GraphQL query string. 

67 description: Optional description. 

68 operation_name: Optional operation name. 

69 

70 Returns: 

71 The created whitelist entry. 

72 """ 

73 hash_value = self._compute_hash(query) 

74 

75 entry = WhitelistEntry( 

76 hash=hash_value, 

77 query=query, 

78 operation_name=operation_name, 

79 description=description, 

80 ) 

81 

82 self._entries[hash_value] = entry 

83 return entry 

84 

85 def remove(self, query: str) -> bool: 

86 """Remove a query from the whitelist. 

87 

88 Args: 

89 query: GraphQL query string. 

90 

91 Returns: 

92 True if removed, False if not found. 

93 """ 

94 hash_value = self._compute_hash(query) 

95 

96 if hash_value in self._entries: 

97 del self._entries[hash_value] 

98 return True 

99 

100 return False 

101 

102 def is_allowed(self, query: str) -> bool: 

103 """Check if a query is allowed. 

104 

105 Args: 

106 query: GraphQL query string. 

107 

108 Returns: 

109 True if allowed, False otherwise. 

110 """ 

111 if not self._enabled: 

112 return True 

113 

114 hash_value = self._compute_hash(query) 

115 return hash_value in self._entries 

116 

117 def get_entry(self, query: str) -> WhitelistEntry | None: 

118 """Get the whitelist entry for a query. 

119 

120 Args: 

121 query: GraphQL query string. 

122 

123 Returns: 

124 Whitelist entry if found, None otherwise. 

125 """ 

126 hash_value = self._compute_hash(query) 

127 return self._entries.get(hash_value) 

128 

129 def get_all_hashes(self) -> set[str]: 

130 """Get all whitelisted query hashes. 

131 

132 Returns: 

133 Set of query hashes. 

134 """ 

135 return set(self._entries.keys()) 

136 

137 def clear(self) -> None: 

138 """Clear all whitelist entries.""" 

139 self._entries.clear() 

140 

141 def _compute_hash(self, query: str) -> str: 

142 """Compute SHA256 hash of a query. 

143 

144 Args: 

145 query: GraphQL query string. 

146 

147 Returns: 

148 Hex-encoded hash. 

149 """ 

150 # Normalize query (remove extra whitespace) 

151 normalized = " ".join(query.split()) 

152 return str(hashing.hash_hex(normalized)) 

153 

154 

155# Default whitelist instance 

156default_whitelist = QueryWhitelist() 

157 

158 

159__all__ = [ 

160 "QueryWhitelist", 

161 "WhitelistEntry", 

162 "default_whitelist", 

163]