Mercurial Repositories > eric / file comparison
comparison: WebBrowser/SafeBrowsing/SafeBrowsingCache.py
WebBrowser/SafeBrowsing/SafeBrowsingCache.py
- branch
- safe_browsing
- changeset 5819
- 69fa45e95673
- parent 5818
- cae9956be67e
- child 5820
- b610cb5b501a
equal
deleted
inserted
replaced
| 5818:cae9956be67e | 5819:69fa45e95673 |
|---|---|
| 15 # | 15 # |
| 16 | 16 |
| 17 from __future__ import unicode_literals | 17 from __future__ import unicode_literals |
| 18 | 18 |
| 19 import os | 19 import os |
| 20 import hashlib | |
| 20 | 21 |
| 21 from PyQt5.QtCore import QObject | 22 from PyQt5.QtCore import QObject |
| 22 from PyQt5.QtSql import QSql, QSqlDatabase, QSqlQuery | 23 from PyQt5.QtSql import QSql, QSqlDatabase, QSqlQuery |
| 24 | |
| 25 from .SafeBrowsingUtilities import toHex | |
| 23 | 26 |
| 24 | 27 |
| 25 class ThreatList(object): | 28 class ThreatList(object): |
| 26 """ | 29 """ |
| 27 Class implementing the threat list info. | 30 Class implementing the threat list info. |
| 159 self.__dbFileName = os.path.join(dbPath, "SafeBrowsingCache.db") | 162 self.__dbFileName = os.path.join(dbPath, "SafeBrowsingCache.db") |
| 160 preparationNeeded = not os.path.exists(self.__dbFileName) | 163 preparationNeeded = not os.path.exists(self.__dbFileName) |
| 161 | 164 |
| 162 self.__openCacheDb() | 165 self.__openCacheDb() |
| 163 if preparationNeeded: | 166 if preparationNeeded: |
| 164 self.__prepareCacheDb() | 167 self.prepareCacheDb() |
| 165 | 168 |
| 166 def close(self): | 169 def close(self): |
| 167 """ | 170 """ |
| 168 Public method to close the database. | 171 Public method to close the database. |
| 169 """ | 172 """ |
| 188 QSqlDatabase.removeDatabase(self.__connectionName) | 191 QSqlDatabase.removeDatabase(self.__connectionName) |
| 189 else: | 192 else: |
| 190 opened = True | 193 opened = True |
| 191 return opened | 194 return opened |
| 192 | 195 |
| 193 def __prepareCacheDb(self): | 196 def prepareCacheDb(self): |
| 194 """ | 197 """ |
| 195 Private method to prepare the cache database. | 198 Public method to prepare the cache database. |
| 196 """ | 199 """ |
| 197 db = QSqlDatabase.database(self.__connectionName) | 200 db = QSqlDatabase.database(self.__connectionName) |
| 198 db.transaction() | 201 db.transaction() |
| 199 try: | 202 try: |
| 200 query = QSqlQuery(db) | 203 query = QSqlQuery(db) |
| 330 updateQueryStr = """ | 333 updateQueryStr = """ |
| 331 UPDATE full_hash SET | 334 UPDATE full_hash SET |
| 332 expires_at=datetime(current_timestamp, '+{0} SECONDS') | 335 expires_at=datetime(current_timestamp, '+{0} SECONDS') |
| 333 WHERE value=? AND threat_type=? AND platform_type=? AND | 336 WHERE value=? AND threat_type=? AND platform_type=? AND |
| 334 threat_entry_type=? | 337 threat_entry_type=? |
| 335 """ | 338 """.format(int(cacheDuration)) |
| 336 | 339 |
| 337 db = QSqlDatabase.database(self.__connectionName) | 340 db = QSqlDatabase.database(self.__connectionName) |
| 338 if db.isOpen(): | 341 if db.isOpen(): |
| 339 db.transaction() | 342 db.transaction() |
| 340 try: | 343 try: |
| 347 query.addBindValue(malwareThreatType) | 350 query.addBindValue(malwareThreatType) |
| 348 query.exec_() | 351 query.exec_() |
| 349 del query | 352 del query |
| 350 | 353 |
| 351 query = QSqlQuery(db) | 354 query = QSqlQuery(db) |
| 352 query.prepare(updateQueryStr.format(int(cacheDuration))) | 355 query.prepare(updateQueryStr) |
| 353 query.addBindValue(hashValue, QSql.In | QSql.Binary) | 356 query.addBindValue(hashValue, QSql.In | QSql.Binary) |
| 354 query.addBindValue(threatList.threatType) | 357 query.addBindValue(threatList.threatType) |
| 355 query.addBindValue(threatList.platformType) | 358 query.addBindValue(threatList.platformType) |
| 356 query.addBindValue(threatList.threatEntryType) | 359 query.addBindValue(threatList.threatEntryType) |
| 357 query.exec_() | 360 query.exec_() |
| 394 @type int or float | 397 @type int or float |
| 395 """ | 398 """ |
| 396 queryStr = """ | 399 queryStr = """ |
| 397 DELETE FROM full_hash | 400 DELETE FROM full_hash |
| 398 WHERE expires_at=datetime(current_timestamp, '{0} SECONDS') | 401 WHERE expires_at=datetime(current_timestamp, '{0} SECONDS') |
| 399 """ | 402 """.format(int(keepExpiredFor)) |
| 400 | 403 |
| 401 db = QSqlDatabase.database(self.__connectionName) | 404 db = QSqlDatabase.database(self.__connectionName) |
| 402 if db.isOpen(): | 405 if db.isOpen(): |
| 403 db.transaction() | 406 db.transaction() |
| 404 try: | 407 try: |
| 405 query = QSqlQuery(db) | 408 query = QSqlQuery(db) |
| 406 query.prepare(queryStr.format(int(keepExpiredFor))) | 409 query.prepare(queryStr) |
| 407 query.exec_() | 410 query.exec_() |
| 408 del query | 411 del query |
| 409 finally: | 412 finally: |
| 410 db.commit() | 413 db.commit() |
| 414 | |
| 415 def updateHashPrefixExpiration(self, threatList, hashPrefix, | |
| 416 negativeCacheDuration): | |
| 417 """ | |
| 418 Public method to update the hash prefix expiration time. | |
| 419 | |
| 420 @param threatList threat list info object | |
| 421 @type ThreatList | |
| 422 @param hashPrefix hash prefix | |
| 423 @type bytes | |
| 424 @param negativeCacheDuration time in seconds the entry should remain | |
| 425 in the cache | |
| 426 @type int or float | |
| 427 """ | |
| 428 queryStr = """ | |
| 429 UPDATE hash_prefix | |
| 430 SET negative_expires_at=datetime(current_timestamp, '+{0} SECONDS') | |
| 431 WHERE value=? AND threat_type=? AND platform_type=? AND | |
| 432 threat_entry_type=? | |
| 433 """.format(int(negativeCacheDuration)) | |
| 434 | |
| 435 db = QSqlDatabase.database(self.__connectionName) | |
| 436 if db.isOpen(): | |
| 437 db.transaction() | |
| 438 try: | |
| 439 query = QSqlQuery(db) | |
| 440 query.prepare(queryStr) | |
| 441 query.addBindValue(hashPrefix, QSql.In | QSql.Binary) | |
| 442 query.addBindValue(threatList.threatType) | |
| 443 query.addBindValue(threatList.platformType) | |
| 444 query.addBindValue(threatList.threatEntryType) | |
| 445 query.exec_() | |
| 446 del query | |
| 447 finally: | |
| 448 db.commit() | |
| 449 | |
| 450 def getThreatLists(self): | |
| 451 """ | |
| 452 Public method to get the available threat lists. | |
| 453 | |
| 454 @return list of available threat lists | |
| 455 @rtype list of tuples of (ThreatList, str) | |
| 456 """ | |
| 457 queryStr = """ | |
| 458 SELECT threat_type,platform_type,threat_entry_type,client_state | |
| 459 FROM threat_list | |
| 460 """ | |
| 461 output = [] | |
| 462 | |
| 463 db = QSqlDatabase.database(self.__connectionName) | |
| 464 if db.isOpen(): | |
| 465 db.transaction() | |
| 466 try: | |
| 467 query = QSqlQuery(db) | |
| 468 query.prepare(queryStr) | |
| 469 | |
| 470 query.exec_() | |
| 471 | |
| 472 while query.next(): | |
| 473 threatType = query.value(0) | |
| 474 platformType = query.value(1) | |
| 475 threatEntryType = query.value(2) | |
| 476 clientState = query.value(3) | |
| 477 threatList = ThreatList(threatType, platformType, | |
| 478 threatEntryType) | |
| 479 output.append((threatList, clientState)) | |
| 480 del query | |
| 481 finally: | |
| 482 db.commit() | |
| 483 | |
| 484 return output | |
| 485 | |
| 486 def addThreatList(self, threatList): | |
| 487 """ | |
| 488 Public method to add a threat list to the cache. | |
| 489 | |
| 490 @param threatList threat list to be added | |
| 491 @type ThreatList | |
| 492 """ | |
| 493 queryStr = """ | |
| 494 INSERT OR IGNORE INTO threat_list | |
| 495 (threat_type, platform_type, threat_entry_type, timestamp) | |
| 496 VALUES (?, ?, ?, current_timestamp) | |
| 497 """ | |
| 498 | |
| 499 db = QSqlDatabase.database(self.__connectionName) | |
| 500 if db.isOpen(): | |
| 501 db.transaction() | |
| 502 try: | |
| 503 query = QSqlQuery(db) | |
| 504 query.prepare(queryStr) | |
| 505 query.addBindValue(threatList.threatType) | |
| 506 query.addBindValue(threatList.platformType) | |
| 507 query.addBindValue(threatList.threatEntryType) | |
| 508 query.exec_() | |
| 509 del query | |
| 510 finally: | |
| 511 db.commit() | |
| 512 | |
| 513 def deleteThreatList(self, threatList): | |
| 514 """ | |
| 515 Public method to delete a threat list from the cache. | |
| 516 | |
| 517 @param threatlist threat list to be deleted | |
| 518 @type ThreatList | |
| 519 """ | |
| 520 queryStr = """ | |
| 521 DELETE FROM threat_list | |
| 522 WHERE threat_type=? AND platform_type=? AND threat_entry_type=? | |
| 523 """ | |
| 524 | |
| 525 db = QSqlDatabase.database(self.__connectionName) | |
| 526 if db.isOpen(): | |
| 527 db.transaction() | |
| 528 try: | |
| 529 query = QSqlQuery(db) | |
| 530 query.prepare(queryStr) | |
| 531 query.addBindValue(threatList.threatType) | |
| 532 query.addBindValue(threatList.platformType) | |
| 533 query.addBindValue(threatList.threatEntryType) | |
| 534 query.exec_() | |
| 535 del query | |
| 536 finally: | |
| 537 db.commit() | |
| 538 | |
| 539 def updateThreatListClientState(self, threatList, clientState): | |
| 540 """ | |
| 541 Public method to update the client state of a threat list. | |
| 542 | |
| 543 @param threatList threat list to update the client state for | |
| 544 @type ThreatList | |
| 545 @param clientState new client state | |
| 546 @type str | |
| 547 """ | |
| 548 queryStr = """ | |
| 549 UPDATE threat_list SET timestamp=current_timestamp, client_state=? | |
| 550 WHERE threat_type=? AND platform_type=? AND threat_entry_type=? | |
| 551 """ | |
| 552 | |
| 553 db = QSqlDatabase.database(self.__connectionName) | |
| 554 if db.isOpen(): | |
| 555 db.transaction() | |
| 556 try: | |
| 557 query = QSqlQuery(db) | |
| 558 query.prepare(queryStr) | |
| 559 query.addBindValue(clientState) | |
| 560 query.addBindValue(threatList.threatType) | |
| 561 query.addBindValue(threatList.platformType) | |
| 562 query.addBindValue(threatList.threatEntryType) | |
| 563 query.exec_() | |
| 564 del query | |
| 565 finally: | |
| 566 db.commit() | |
| 567 | |
| 568 def hashPrefixListChecksum(self, threatList): | |
| 569 """ | |
| 570 Public method to calculate the SHA256 checksum for an alphabetically | |
| 571 sorted concatenated list of hash prefixes. | |
| 572 | |
| 573 @param threatList threat list to calculate checksum for | |
| 574 @type ThreatList | |
| 575 @return SHA256 checksum | |
| 576 @rtype bytes | |
| 577 """ | |
| 578 queryStr = """ | |
| 579 SELECT value FROM hash_prefix | |
| 580 WHERE threat_type=? AND platform_type=? AND threat_entry_type=? | |
| 581 ORDER BY value | |
| 582 """ | |
| 583 checksum = None | |
| 584 | |
| 585 db = QSqlDatabase.database(self.__connectionName) | |
| 586 if db.isOpen(): | |
| 587 db.transaction() | |
| 588 allHashes = b"" | |
| 589 try: | |
| 590 query = QSqlQuery(db) | |
| 591 query.prepare(queryStr) | |
| 592 query.addBindValue(threatList.threatType) | |
| 593 query.addBindValue(threatList.platformType) | |
| 594 query.addBindValue(threatList.threatEntryType) | |
| 595 | |
| 596 query.exec_() | |
| 597 | |
| 598 while query.next(): | |
| 599 allHashes += bytes(query.value(0)) | |
| 600 del query | |
| 601 finally: | |
| 602 db.commit() | |
| 603 | |
| 604 checksum = hashlib.sha256(allHashes).digest() | |
| 605 | |
| 606 return checksum | |
| 607 | |
| 608 def populateHashPrefixList(self, threatList, prefixes): | |
| 609 """ | |
| 610 Public method to populate the hash prefixes for a threat list. | |
| 611 | |
| 612 @param threatList threat list of the hash prefixes | |
| 613 @type ThreatList | |
| 614 @param prefixes hash prefixes to be inserted | |
| 615 @type bytes | |
| 616 """ | |
| 617 queryStr = """ | |
| 618 INSERT INTO hash_prefix | |
| 619 (value, cue, threat_type, platform_type, threat_entry_type, | |
| 620 timestamp) | |
| 621 VALUES (?, ?, ?, ?, ?, current_timestamp) | |
| 622 """ | |
| 623 | |
| 624 db = QSqlDatabase.database(self.__connectionName) | |
| 625 if db.isOpen(): | |
| 626 db.transaction() | |
| 627 try: | |
| 628 for prefix in prefixes: | |
| 629 query = QSqlQuery(db) | |
| 630 query.prepare(queryStr) | |
| 631 query.addBindValue(prefix, QSql.In | QSql.Binary) | |
| 632 query.addBindValue(toHex(prefix[:4])) | |
| 633 query.addBindValue(threatList.threatType) | |
| 634 query.addBindValue(threatList.platformType) | |
| 635 query.addBindValue(threatList.threatEntryType) | |
| 636 query.exec_() | |
| 637 del query | |
| 638 finally: | |
| 639 db.commit() | |
| 640 | |
| 641 def getHashPrefixValuesToRemove(self, threatList, indexes): | |
| 642 """ | |
| 643 Public method to get the hash prefix values to be removed from the | |
| 644 cache. | |
| 645 | |
| 646 @param threatList threat list to remove prefixes from | |
| 647 @type ThreatList | |
| 648 @param indexes list of indexes of prefixes to be removed | |
| 649 @type list of int | |
| 650 @return list of hash prefixes to be removed | |
| 651 @rtype list of bytes | |
| 652 """ | |
| 653 queryStr = """ | |
| 654 SELECT value FROM hash_prefix | |
| 655 WHERE threat_type=? AND platform_type=? AND threat_entry_type=? | |
| 656 ORDER BY value | |
| 657 """ | |
| 658 indexes = set(indexes) | |
| 659 output = [] | |
| 660 | |
| 661 db = QSqlDatabase.database(self.__connectionName) | |
| 662 if db.isOpen(): | |
| 663 db.transaction() | |
| 664 try: | |
| 665 query = QSqlQuery(db) | |
| 666 query.prepare(queryStr) | |
| 667 query.addBindValue(threatList.threatType) | |
| 668 query.addBindValue(threatList.platformType) | |
| 669 query.addBindValue(threatList.threatEntryType) | |
| 670 | |
| 671 query.exec_() | |
| 672 | |
| 673 index = 0 | |
| 674 while query.next(): | |
| 675 if index in indexes: | |
| 676 prefix = bytes(query.value(0)) | |
| 677 output.append(prefix) | |
| 678 index += 1 | |
| 679 del query | |
| 680 finally: | |
| 681 db.commit() | |
| 682 | |
| 683 return output | |
| 684 | |
| 685 def removeHashPrefixIndices(self, threatList, indexes): | |
| 686 """ | |
| 687 Public method to remove hash prefixes from the cache. | |
| 688 | |
| 689 @param threatList threat list to delete hash prefixes of | |
| 690 @type ThreatList | |
| 691 @param indexes list of indexes of prefixes to be removed | |
| 692 @type list of int | |
| 693 """ | |
| 694 queryStr = """ | |
| 695 DELETE FROM hash_prefix | |
| 696 WHERE threat_type=? AND platform_type=? AND | |
| 697 threat_entry_type=? AND value IN ({0}) | |
| 698 """ | |
| 699 batchSize = 40 | |
| 700 | |
| 701 prefixesToRemove = self.getHashPrefixValuesToRemove( | |
| 702 threatList, indexes) | |
| 703 if prefixesToRemove: | |
| 704 db = QSqlDatabase.database(self.__connectionName) | |
| 705 if db.isOpen(): | |
| 706 db.transaction() | |
| 707 try: | |
| 708 for index in range(0, len(prefixesToRemove), batchSize): | |
| 709 removeBatch = \ | |
| 710 prefixesToRemove[index:(index + batchSize)] | |
| 711 | |
| 712 query = QSqlQuery(db) | |
| 713 query.prepare( | |
| 714 queryStr.format(",".join(["?"] * len(removeBatch))) | |
| 715 ) | |
| 716 query.addBindValue(threatList.threatType) | |
| 717 query.addBindValue(threatList.platformType) | |
| 718 query.addBindValue(threatList.threatEntryType) | |
| 719 for prefix in removeBatch: | |
| 720 query.addBindValue(prefix, QSql.In | QSql.Binary) | |
| 721 query.exec_() | |
| 722 del query | |
| 723 finally: | |
| 724 db.commit() |
