Helper class to track stolen files and avoid re-downloading.
| 50 | |
| 51 | |
| 52 | class FileTracker: |
| 53 | """Helper class to track stolen files and avoid re-downloading.""" |
| 54 | |
| 55 | def __init__(self, protocol, datadir): |
| 56 | """ |
| 57 | Initialize the file tracker. |
| 58 | |
| 59 | Args: |
| 60 | protocol: Protocol name (ssh, ftp, smb, rdp, telnet) |
| 61 | datadir: Data directory for storing tracking files |
| 62 | """ |
| 63 | self.protocol = protocol |
| 64 | self.db_file = os.path.join(datadir, f'stolen_files_{protocol}.json') |
| 65 | self.stolen_files = {} |
| 66 | self._load_db() |
| 67 | |
| 68 | def _load_db(self): |
| 69 | """Load the database of already stolen files.""" |
| 70 | try: |
| 71 | if os.path.exists(self.db_file): |
| 72 | with open(self.db_file, 'r') as f: |
| 73 | self.stolen_files = json.load(f) |
| 74 | else: |
| 75 | self.stolen_files = {} |
| 76 | except Exception as e: |
| 77 | logger.warning(f"Could not load stolen files database for {self.protocol}: {e}") |
| 78 | self.stolen_files = {} |
| 79 | |
| 80 | def _save_db(self): |
| 81 | """Save the database of stolen files.""" |
| 82 | try: |
| 83 | with open(self.db_file, 'w') as f: |
| 84 | json.dump(self.stolen_files, f, indent=2) |
| 85 | except Exception as e: |
| 86 | logger.error(f"Could not save stolen files database for {self.protocol}: {e}") |
| 87 | |
| 88 | def is_file_stolen(self, ip, remote_file): |
| 89 | """ |
| 90 | Check if a file has already been stolen from this host. |
| 91 | |
| 92 | Args: |
| 93 | ip: IP address of the host |
| 94 | remote_file: Path to the remote file |
| 95 | |
| 96 | Returns: |
| 97 | True if file was already stolen, False otherwise |
| 98 | """ |
| 99 | if ip not in self.stolen_files: |
| 100 | return False |
| 101 | return remote_file in self.stolen_files[ip] |
| 102 | |
| 103 | def mark_file_stolen(self, ip, remote_file): |
| 104 | """ |
| 105 | Mark a file as stolen. |
| 106 | |
| 107 | Args: |
| 108 | ip: IP address of the host |
| 109 | remote_file: Path to the remote file |