| 45 | |
| 46 | |
| 47 | class DataService(DataServiceInterface, BaseService): |
| 48 | |
| 49 | def __init__(self): |
| 50 | self.log = self.add_service('data_svc', self) |
| 51 | self.schema = dict(agents=[], planners=[], adversaries=[], abilities=[], sources=[], operations=[], |
| 52 | schedules=[], plugins=[], obfuscators=[], objectives=[], data_encoders=[]) |
| 53 | self.ram = copy.deepcopy(self.schema) |
| 54 | |
| 55 | @staticmethod |
| 56 | def _iter_data_files(): |
| 57 | """Yield paths to data files managed by caldera. |
| 58 | |
| 59 | The files paths are relative to the root caldera folder, so they |
| 60 | will begin with "data/". |
| 61 | |
| 62 | Note: |
| 63 | This will skip any files starting with '.' (e.g., '.gitkeep'). |
| 64 | """ |
| 65 | for data_glob in DATA_FILE_GLOBS: |
| 66 | for f in glob.glob(data_glob): |
| 67 | yield f |
| 68 | |
| 69 | @staticmethod |
| 70 | def _delete_file(path): |
| 71 | if not os.path.exists(path): |
| 72 | return |
| 73 | elif os.path.isdir(path): |
| 74 | shutil.rmtree(path) |
| 75 | else: |
| 76 | os.remove(path) |
| 77 | |
| 78 | @staticmethod |
| 79 | async def destroy(): |
| 80 | """Reset the caldera data directory and server state. |
| 81 | |
| 82 | This creates a gzipped tarball backup of the data files tracked by caldera. |
| 83 | Paths are preserved within the tarball, with all files having "data/" as the |
| 84 | root. |
| 85 | """ |
| 86 | if not os.path.exists(DATA_BACKUP_DIR): |
| 87 | os.mkdir(DATA_BACKUP_DIR) |
| 88 | |
| 89 | timestamp = datetime.datetime.utcnow().strftime('%Y%m%d%H%M%S') |
| 90 | tarball_path = os.path.join(DATA_BACKUP_DIR, f'backup-{timestamp}.tar.gz') |
| 91 | |
| 92 | with tarfile.open(tarball_path, 'w:gz') as tarball: |
| 93 | for file_path in DataService._iter_data_files(): |
| 94 | tarball.add(file_path) |
| 95 | DataService._delete_file(file_path) |
| 96 | |
| 97 | async def save_state(self): |
| 98 | await self._prune_non_critical_data() |
| 99 | await self.get_service('file_svc').save_file('object_store', pickle.dumps(self.ram), 'data') |
| 100 | |
| 101 | async def restore_state(self): |
| 102 | """ |
| 103 | Restore the object database |
| 104 |
no outgoing calls