mirror of https://github.com/CIRCL/lookyloo
128 lines
3.8 KiB
Python
Executable File
128 lines
3.8 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
from subprocess import Popen
|
|
|
|
from redis import Redis
|
|
from redis.exceptions import ConnectionError
|
|
|
|
from lookyloo.default import get_homedir, get_socket_path, get_config
|
|
|
|
|
|
def check_running(name: str) -> bool:
|
|
socket_path = get_socket_path(name)
|
|
if not os.path.exists(socket_path):
|
|
return False
|
|
try:
|
|
r = Redis(unix_socket_path=socket_path)
|
|
return True if r.ping() else False
|
|
except ConnectionError:
|
|
return False
|
|
|
|
|
|
def launch_cache(storage_directory: Path | None=None) -> None:
|
|
if not storage_directory:
|
|
storage_directory = get_homedir()
|
|
if not check_running('cache'):
|
|
Popen(["./run_redis.sh"], cwd=(storage_directory / 'cache'))
|
|
|
|
|
|
def shutdown_cache(storage_directory: Path | None=None) -> None:
|
|
if not storage_directory:
|
|
storage_directory = get_homedir()
|
|
r = Redis(unix_socket_path=get_socket_path('cache'))
|
|
r.shutdown(save=True)
|
|
print('Redis cache database shutdown.')
|
|
|
|
|
|
def launch_indexing(storage_directory: Path | None=None) -> None:
|
|
if not storage_directory:
|
|
storage_directory = get_homedir()
|
|
if not check_running('indexing'):
|
|
Popen(["./run_redis.sh"], cwd=(storage_directory / 'indexing'))
|
|
|
|
|
|
def shutdown_indexing(storage_directory: Path | None=None) -> None:
|
|
if not storage_directory:
|
|
storage_directory = get_homedir()
|
|
r = Redis(unix_socket_path=get_socket_path('indexing'))
|
|
r.shutdown(save=True)
|
|
print('Redis indexing database shutdown.')
|
|
|
|
|
|
def launch_full_index(storage_directory: Path | None=None) -> None:
|
|
if not storage_directory:
|
|
storage_directory = get_homedir()
|
|
if not check_running('full_index'):
|
|
Popen(["./run_kvrocks.sh"], cwd=(storage_directory / 'full_index'))
|
|
|
|
|
|
def shutdown_full_index(storage_directory: Path | None=None) -> None:
|
|
if not storage_directory:
|
|
storage_directory = get_homedir()
|
|
r = Redis(unix_socket_path=get_socket_path('full_index'))
|
|
r.shutdown()
|
|
print('Kvrocks full indexing database shutdown.')
|
|
|
|
|
|
def launch_all() -> None:
|
|
launch_cache()
|
|
launch_indexing()
|
|
if get_config('generic', 'index_everything'):
|
|
launch_full_index()
|
|
|
|
|
|
def check_all(stop: bool=False) -> None:
|
|
backends: dict[str, bool] = {'cache': False, 'indexing': False}
|
|
if get_config('generic', 'index_everything'):
|
|
backends['full_index'] = False
|
|
while True:
|
|
for db_name in backends.keys():
|
|
try:
|
|
backends[db_name] = check_running(db_name)
|
|
except Exception:
|
|
backends[db_name] = False
|
|
if stop:
|
|
if not any(running for running in backends.values()):
|
|
break
|
|
else:
|
|
if all(running for running in backends.values()):
|
|
break
|
|
for db_name, running in backends.items():
|
|
if not stop and not running:
|
|
print(f"Waiting on {db_name} to start")
|
|
if stop and running:
|
|
print(f"Waiting on {db_name} to stop")
|
|
time.sleep(1)
|
|
|
|
|
|
def stop_all() -> None:
|
|
shutdown_cache()
|
|
shutdown_indexing()
|
|
if get_config('generic', 'index_everything'):
|
|
shutdown_full_index()
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description='Manage backend DBs.')
|
|
parser.add_argument("--start", action='store_true', default=False, help="Start all")
|
|
parser.add_argument("--stop", action='store_true', default=False, help="Stop all")
|
|
parser.add_argument("--status", action='store_true', default=True, help="Show status")
|
|
args = parser.parse_args()
|
|
|
|
if args.start:
|
|
launch_all()
|
|
if args.stop:
|
|
stop_all()
|
|
if not args.stop and args.status:
|
|
check_all()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|