78 lines
3.0 KiB
Python
78 lines
3.0 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
# -*- coding: utf-8 -*-
|
||
|
|
|
||
|
|
import io
|
||
|
|
import re
|
||
|
|
import shutil
|
||
|
|
from datetime import datetime
|
||
|
|
|
||
|
|
import psutil
|
||
|
|
import pyudev
|
||
|
|
from flask import jsonify, send_file
|
||
|
|
|
||
|
|
|
||
|
|
class Save():
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self.mount_point = ""
|
||
|
|
return None
|
||
|
|
|
||
|
|
def usb_check(self) -> dict:
|
||
|
|
"""Check if an USB storage is connected or not.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
dict: contains the connection status.
|
||
|
|
"""
|
||
|
|
self.usb_devices = []
|
||
|
|
context = pyudev.Context()
|
||
|
|
removable = [device for device in context.list_devices(
|
||
|
|
subsystem='block', DEVTYPE='disk')]
|
||
|
|
for device in removable:
|
||
|
|
if "usb" in device.sys_path:
|
||
|
|
partitions = [device.device_node for device in context.list_devices(
|
||
|
|
subsystem='block', DEVTYPE='partition', parent=device)]
|
||
|
|
for p in psutil.disk_partitions():
|
||
|
|
if p.device in partitions:
|
||
|
|
self.mount_point = p.mountpoint
|
||
|
|
return jsonify({"status": True,
|
||
|
|
"message": "USB storage connected"})
|
||
|
|
self.mount_point = ""
|
||
|
|
return jsonify({"status": False,
|
||
|
|
"message": "USB storage not connected"})
|
||
|
|
|
||
|
|
def save_capture(self, token, method) -> any:
|
||
|
|
"""Save the capture to the USB device or push a ZIP
|
||
|
|
file to download.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
token (str): capture token
|
||
|
|
method (str): method used to save
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
dict: operation status OR Flask answer.
|
||
|
|
"""
|
||
|
|
if re.match(r"[A-F0-9]{8}", token):
|
||
|
|
try:
|
||
|
|
if method == "usb":
|
||
|
|
cd = datetime.now().strftime("%d%m%Y-%H%M")
|
||
|
|
if shutil.make_archive("{}/SpyGuard_{}".format(self.mount_point, cd), "zip", "/tmp/{}/".format(token)):
|
||
|
|
shutil.rmtree("/tmp/{}/".format(token))
|
||
|
|
return jsonify({"status": True,
|
||
|
|
"message": "Capture saved on the USB key"})
|
||
|
|
elif method == "url":
|
||
|
|
cd = datetime.now().strftime("%d%m%Y-%H%M")
|
||
|
|
if shutil.make_archive("/tmp/SpyGuard_{}".format(cd), "zip", "/tmp/{}/".format(token)):
|
||
|
|
shutil.rmtree("/tmp/{}/".format(token))
|
||
|
|
with open("/tmp/SpyGuard_{}.zip".format(cd), "rb") as f:
|
||
|
|
return send_file(
|
||
|
|
io.BytesIO(f.read()),
|
||
|
|
mimetype="application/octet-stream",
|
||
|
|
as_attachment=True,
|
||
|
|
attachment_filename="SpyGuard_{}.zip".format(cd))
|
||
|
|
except:
|
||
|
|
return jsonify({"status": False,
|
||
|
|
"message": "Error while saving capture"})
|
||
|
|
else:
|
||
|
|
return jsonify({"status": False,
|
||
|
|
"message": "Bad token value"})
|