Coverage for ivatar/test_file_security.py: 98%

1# -*- coding: utf-8 -*- 

2""" 

3Tests for file upload security enhancements 

4""" 

5 

6from unittest.mock import patch 

7 

8from django.test import TestCase, override_settings 

9from django.core.files.uploadedfile import SimpleUploadedFile 

10from django.contrib.auth.models import User 

11 

12from ivatar.file_security import ( 

13 FileValidator, 

14 validate_uploaded_file, 

15 get_file_security_report, 

16) 

17from ivatar.ivataraccount.forms import UploadPhotoForm 

18 

19 

20class FileSecurityTestCase(TestCase): 

21 """Test cases for file upload security""" 

22 

23 def setUp(self): 

24 """Set up test data""" 

25 self.user = User.objects.create_user( 

26 username="testuser", email="test@example.com", password="testpass123" 

27 ) 

28 

29 # Create test image data 

30 self.valid_jpeg_data = b"\xff\xd8\xff\xe0\x00\x10JFIF\x00\x01\x01\x01\x00H\x00H\x00\x00\xff\xdb\x00C\x00\x08\x06\x06\x07\x06\x05\x08\x07\x07\x07\t\t\x08\n\x0c\x14\r\x0c\x0b\x0b\x0c\x19\x12\x13\x0f\x14\x1d\x1a\x1f\x1e\x1d\x1a\x1c\x1c $.' \",#\x1c\x1c(7),01444\x1f'9=82<.342\xff\xc0\x00\x11\x08\x00\x01\x00\x01\x01\x01\x11\x00\x02\x11\x01\x03\x11\x01\xff\xc4\x00\x1f\x00\x00\x01\x05\x01\x01\x01\x01\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x01\x02\x03\x04\x05\x06\x07\x08\t\n\x0b\xff\xc4\x00\xb5\x10\x00\x02\x01\x03\x03\x02\x04\x03\x05\x05\x04\x04\x00\x00\x01}\x01\x02\x03\x00\x04\x11\x05\x12!1A\x06\x13Qa\x07\"q\x142\x81\x91\xa1\x08#B\xb1\xc1\x15R\xd1\xf0$3br\x82\t\n\x16\x17\x18\x19\x1a%&'()*456789:CDEFGHIJSTUVWXYZcdefghijstuvwxyz\x83\x84\x85\x86\x87\x88\x89\x8a\x92\x93\x94\x95\x96\x97\x98\x99\x9a\xa2\xa3\xa4\xa5\xa6\xa7\xa8\xa9\xaa\xb2\xb3\xb4\xb5\xb6\xb7\xb8\xb9\xba\xc2\xc3\xc4\xc5\xc6\xc7\xc8\xc9\xca\xd2\xd3\xd4\xd5\xd6\xd7\xd8\xd9\xda\xe1\xe2\xe3\xe4\xe5\xe6\xe7\xe8\xe9\xea\xf1\xf2\xf3\xf4\xf5\xf6\xf7\xf8\xf9\xfa\xff\xda\x00\x0c\x03\x01\x00\x02\x11\x03\x11\x00\x3f\x00\xf9\xff\xd9" 

31 

32 self.malicious_data = b'GIF89a<script>alert("xss")</script>' 

33 self.large_data = b"x" * (10 * 1024 * 1024) # 10MB 

34 

35 def tearDown(self): 

36 """Clean up after tests""" 

37 pass 

38 

39 def test_valid_jpeg_validation(self): 

40 """Test validation of valid JPEG file""" 

41 validator = FileValidator(self.valid_jpeg_data, "test.jpg") 

42 

43 # Mock PIL validation to avoid issues with test data 

44 with patch.object(validator, "validate_pil_image") as mock_pil: 

45 mock_pil.return_value = { 

46 "valid": True, 

47 "image_info": { 

48 "format": "JPEG", 

49 "mode": "RGB", 

50 "size": (100, 100), 

51 "width": 100, 

52 "height": 100, 

53 "has_transparency": False, 

54 }, 

55 "errors": [], 

56 "warnings": [], 

57 } 

58 

59 results = validator.comprehensive_validation() 

60 

61 self.assertTrue(results["valid"]) 

62 self.assertEqual(results["file_info"]["detected_type"], "image/jpeg") 

63 self.assertGreaterEqual(results["security_score"], 80) 

64 

65 def test_magic_bytes_validation(self): 

66 """Test magic bytes validation""" 

67 validator = FileValidator(self.valid_jpeg_data, "test.jpg") 

68 results = validator.validate_magic_bytes() 

69 

70 self.assertTrue(results["valid"]) 

71 self.assertEqual(results["detected_type"], "image/jpeg") 

72 

73 def test_malicious_content_detection(self): 

74 """Test detection of malicious content""" 

75 validator = FileValidator(self.malicious_data, "malicious.gif") 

76 results = validator.scan_for_malicious_content() 

77 

78 self.assertTrue(results["suspicious"]) 

79 self.assertGreater(len(results["threats"]), 0) 

80 

81 def test_file_size_validation(self): 

82 """Test file size validation""" 

83 validator = FileValidator(self.large_data, "large.jpg") 

84 results = validator.validate_basic() 

85 

86 self.assertFalse(results["valid"]) 

87 self.assertIn("File too large", results["errors"][0]) 

88 

89 def test_invalid_extension_validation(self): 

90 """Test invalid file extension validation""" 

91 validator = FileValidator(self.valid_jpeg_data, "test.exe") 

92 results = validator.validate_basic() 

93 

94 self.assertFalse(results["valid"]) 

95 self.assertIn("File extension not allowed", results["errors"][0]) 

96 

97 def test_exif_sanitization(self): 

98 """Test EXIF data sanitization""" 

99 validator = FileValidator(self.valid_jpeg_data, "test.jpg") 

100 sanitized_data = validator.sanitize_exif_data() 

101 

102 # Should return data (may be same or sanitized) 

103 self.assertIsInstance(sanitized_data, bytes) 

104 self.assertGreater(len(sanitized_data), 0) 

105 

106 def test_comprehensive_validation_function(self): 

107 """Test the main validation function""" 

108 # Mock PIL validation to avoid issues with test data 

109 with patch("ivatar.file_security.FileValidator.validate_pil_image") as mock_pil: 

110 mock_pil.return_value = { 

111 "valid": True, 

112 "image_info": {"format": "JPEG", "size": (100, 100)}, 

113 "errors": [], 

114 "warnings": [], 

115 } 

116 

117 is_valid, results, sanitized_data = validate_uploaded_file( 

118 self.valid_jpeg_data, "test.jpg" 

119 ) 

120 

121 self.assertTrue(is_valid) 

122 self.assertIsInstance(results, dict) 

123 self.assertIsInstance(sanitized_data, bytes) 

124 

125 def test_security_report_generation(self): 

126 """Test security report generation""" 

127 # Mock PIL validation to avoid issues with test data 

128 with patch("ivatar.file_security.FileValidator.validate_pil_image") as mock_pil: 

129 mock_pil.return_value = { 

130 "valid": True, 

131 "image_info": {"format": "JPEG", "size": (100, 100)}, 

132 "errors": [], 

133 "warnings": [], 

134 } 

135 

136 report = get_file_security_report(self.valid_jpeg_data, "test.jpg") 

137 

138 self.assertIn("valid", report) 

139 self.assertIn("security_score", report) 

140 self.assertIn("file_info", report) 

141 

142 @patch("ivatar.file_security.magic.from_buffer") 

143 def test_mime_type_validation(self, mock_magic): 

144 """Test MIME type validation with mocked magic""" 

145 mock_magic.return_value = "image/jpeg" 

146 

147 validator = FileValidator(self.valid_jpeg_data, "test.jpg") 

148 results = validator.validate_mime_type() 

149 

150 self.assertTrue(results["valid"]) 

151 self.assertEqual(results["detected_mime"], "image/jpeg") 

152 

153 def test_polyglot_attack_detection(self): 

154 """Test detection of polyglot attacks""" 

155 polyglot_data = b'GIF89a<script>alert("xss")</script>' 

156 validator = FileValidator(polyglot_data, "polyglot.gif") 

157 results = validator.scan_for_malicious_content() 

158 

159 self.assertTrue(results["suspicious"]) 

160 # Check for either polyglot attack or suspicious script pattern 

161 threats_text = " ".join(results["threats"]).lower() 

162 self.assertTrue( 

163 "polyglot attack" in threats_text or "suspicious pattern" in threats_text, 

164 f"Expected polyglot attack or suspicious pattern, got: {results['threats']}", 

165 ) 

166 

167 

168class UploadPhotoFormSecurityTestCase(TestCase): 

169 """Test cases for UploadPhotoForm security enhancements""" 

170 

171 def setUp(self): 

172 """Set up test data""" 

173 self.user = User.objects.create_user( 

174 username="testuser", email="test@example.com", password="testpass123" 

175 ) 

176 

177 def test_form_validation_with_valid_file(self): 

178 """Test form validation with valid file""" 

179 valid_jpeg_data = b"\xff\xd8\xff\xe0\x00\x10JFIF\x00\x01\x01\x01\x00H\x00H\x00\x00\xff\xdb\x00C\x00\x08\x06\x06\x07\x06\x05\x08\x07\x07\x07\t\t\x08\n\x0c\x14\r\x0c\x0b\x0b\x0c\x19\x12\x13\x0f\x14\x1d\x1a\x1f\x1e\x1d\x1a\x1c\x1c $.' \",#\x1c\x1c(7),01444\x1f'9=82<.342\xff\xc0\x00\x11\x08\x00\x01\x00\x01\x01\x01\x11\x00\x02\x11\x01\x03\x11\x01\xff\xc4\x00\x1f\x00\x00\x01\x05\x01\x01\x01\x01\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x01\x02\x03\x04\x05\x06\x07\x08\t\n\x0b\xff\xc4\x00\xb5\x10\x00\x02\x01\x03\x03\x02\x04\x03\x05\x05\x04\x04\x00\x00\x01}\x01\x02\x03\x00\x04\x11\x05\x12!1A\x06\x13Qa\x07\"q\x142\x81\x91\xa1\x08#B\xb1\xc1\x15R\xd1\xf0$3br\x82\t\n\x16\x17\x18\x19\x1a%&'()*456789:CDEFGHIJSTUVWXYZcdefghijstuvwxyz\x83\x84\x85\x86\x87\x88\x89\x8a\x92\x93\x94\x95\x96\x97\x98\x99\x9a\xa2\xa3\xa4\xa5\xa6\xa7\xa8\xa9\xaa\xb2\xb3\xb4\xb5\xb6\xb7\xb8\xb9\xba\xc2\xc3\xc4\xc5\xc6\xc7\xc8\xc9\xca\xd2\xd3\xd4\xd5\xd6\xd7\xd8\xd9\xda\xe1\xe2\xe3\xe4\xe5\xe6\xe7\xe8\xe9\xea\xf1\xf2\xf3\xf4\xf5\xf6\xf7\xf8\xf9\xfa\xff\xda\x00\x0c\x03\x01\x00\x02\x11\x03\x11\x00\x3f\x00\xf9\xff\xd9" 

180 

181 uploaded_file = SimpleUploadedFile( 

182 "test.jpg", valid_jpeg_data, content_type="image/jpeg" 

183 ) 

184 

185 form_data = {"photo": uploaded_file, "not_porn": True, "can_distribute": True} 

186 

187 form = UploadPhotoForm(data=form_data, files={"photo": uploaded_file}) 

188 

189 # Mock the validation to avoid PIL issues in tests 

190 with patch( 

191 "ivatar.ivataraccount.forms.validate_uploaded_file" 

192 ) as mock_validate: 

193 mock_validate.return_value = ( 

194 True, 

195 {"security_score": 95, "errors": [], "warnings": []}, 

196 valid_jpeg_data, 

197 ) 

198 

199 self.assertTrue(form.is_valid()) 

200 

201 def test_form_validation_with_malicious_file(self): 

202 """Test form validation with malicious file""" 

203 malicious_data = b'GIF89a<script>alert("xss")</script>' 

204 

205 uploaded_file = SimpleUploadedFile( 

206 "malicious.gif", malicious_data, content_type="image/gif" 

207 ) 

208 

209 form_data = {"photo": uploaded_file, "not_porn": True, "can_distribute": True} 

210 

211 form = UploadPhotoForm(data=form_data, files={"photo": uploaded_file}) 

212 

213 # Mock the validation to return malicious file detection 

214 with patch( 

215 "ivatar.ivataraccount.forms.validate_uploaded_file" 

216 ) as mock_validate: 

217 mock_validate.return_value = ( 

218 False, 

219 { 

220 "security_score": 20, 

221 "errors": ["Malicious content detected"], 

222 "warnings": [], 

223 }, 

224 malicious_data, 

225 ) 

226 

227 self.assertFalse(form.is_valid()) 

228 # Check for any error message indicating validation failure 

229 error_text = str(form.errors["photo"]).lower() 

230 self.assertTrue( 

231 "malicious" in error_text or "validation failed" in error_text, 

232 f"Expected malicious or validation failed message, got: {form.errors['photo']}", 

233 ) 

234 

235 

236class UploadPhotoViewSecurityTestCase(TestCase): 

237 """Test cases for UploadPhotoView security enhancements""" 

238 

239 def setUp(self): 

240 """Set up test data""" 

241 self.user = User.objects.create_user( 

242 username="testuser", email="test@example.com", password="testpass123" 

243 ) 

244 

245 def tearDown(self): 

246 """Clean up after tests""" 

247 pass 

248 

249 

250@override_settings( 

251 ENABLE_FILE_SECURITY_VALIDATION=True, 

252 ENABLE_EXIF_SANITIZATION=True, 

253 ENABLE_MALICIOUS_CONTENT_SCAN=True, 

254 ENABLE_RATE_LIMITING=True, 

255) 

256class FileSecurityIntegrationTestCase(TestCase): 

257 """Integration tests for file upload security""" 

258 

259 def setUp(self): 

260 """Set up test data""" 

261 self.user = User.objects.create_user( 

262 username="testuser", email="test@example.com", password="testpass123" 

263 ) 

264 

265 def test_end_to_end_security_validation(self): 

266 """Test end-to-end security validation""" 

267 # This would test the complete flow from upload to storage 

268 # with all security checks enabled 

269 pass 

270 

271 def test_security_logging(self): 

272 """Test that security events are properly logged""" 

273 # This would test that security events are logged 

274 # when malicious files are uploaded 

275 pass