agent-python/htpclient/files.py at 676bf2b7eca31b3f2e5a7444fa3484eae2d2ab89 · hashtopolis/agent-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
import logging
import time
from time import sleep
from pathlib import Path

import os

from htpclient.config import Config
from htpclient.download import Download
from htpclient.helpers import retrieveBinary
from htpclient.initialize import Initialize
from htpclient.jsonRequest import JsonRequest
from htpclient.dicts import *


class Files:
    def __init__(self):
        self.config = Config()
        self.chunk = None
        self.last_check = None
        self.check_interval = 600
        if self.config.get_value('file-deletion-interval'):
            self.check_interval = int(self.config.get_value('file-deletion-interval'))

    def deletion_check(self):
        if self.config.get_value('file-deletion-disable'):
            return
        elif self.last_check is not None and time.time() - self.last_check < self.check_interval:
            return
        query = copy_and_set_token(dict_getFileStatus, self.config.get_value('token'))
        req = JsonRequest(query)
        ans = req.execute()
        self.last_check = time.time()
        if ans is None:
            logging.error("Failed to get file status!")
        elif ans['response'] != 'SUCCESS':
            logging.error("Getting of file status failed: " + str(ans))
        else:
            files = ans['filenames']
            for filename in files:
                file_path = Path(self.config.get_value('files-path'), filename)
                if filename.find("/") != -1 or filename.find("\\") != -1:
                    continue  # ignore invalid file names
                elif os.path.dirname(file_path) != "files":
                    continue  # ignore any case in which we would leave the files folder
                elif os.path.exists(file_path):
                    logging.info("Delete file '" + filename + "' as requested by server...")
                    # When we get the delete requests, this function will check if the <filename>.7z maybe as
                    # an extracted text file. That file will also be deleted.
                    if os.path.splitext(file_path)[1] == '.7z':
                        txt_file = Path(f"{os.path.splitext(file_path)[0]}.txt")
                        if os.path.exists(txt_file):
                            logging.info("Also delete assumed wordlist from archive of same file...")
                            os.unlink(txt_file)
                    os.unlink(file_path)

    def check_files(self, files, task_id):
        for file in files:
            file_localpath = Path(self.config.get_value('files-path'), file)
            txt_file = Path(f"{os.path.splitext(file_localpath)[0]}.txt")
            query = copy_and_set_token(dict_getFile, self.config.get_value('token'))
            query['taskId'] = task_id
            query['file'] = file
            req = JsonRequest(query)
            ans = req.execute()

            # Process request
            if ans is None:
                logging.error("Failed to get file!")
                sleep(5)
                return False
            elif ans['response'] != 'SUCCESS':
                logging.error("Getting of file failed: " + str(ans))
                sleep(5)
                return False
            else:
                # Filesize is OK
                file_size = int(ans['filesize'])
                if os.path.isfile(file_localpath) and os.stat(file_localpath).st_size == file_size:
                    logging.debug("File is present on agent and has matching file size.")
                    continue

                # Multicasting configured
                elif self.config.get_value('multicast'):
                    logging.debug("Multicast is enabled, need to wait until it was delivered!")
                    sleep(5)  # in case the file is not there yet (or not completely), we just wait some time and then try again
                    return False

                # TODO: we might need a better check for this
                if os.path.isfile(txt_file):
                    continue

                # Rsync
                if self.config.get_value('rsync') and Initialize.get_os() != 1:
                    Download.rsync(Path(self.config.get_value('rsync-path'), file), file_localpath)
                else:
                    logging.debug("Starting download of file from server...")
                    Download.download(self.config.get_value('url').replace("api/server.php", "") + ans['url'], file_localpath)

                # Mismatch filesize
                if os.path.isfile(file_localpath) and os.stat(file_localpath).st_size != file_size:
                    logging.error("file size mismatch on file: %s" % file)
                    sleep(5)
                    return False

                # 7z extraction, check if the <filename>.txt does exist.
                if os.path.splitext(file_localpath)[1] == '.7z' and not os.path.isfile(txt_file):
                    # extract if needed
                    files_path = Path(self.config.get_value('files-path'))
                    zr_bin = retrieveBinary("7zr" + Initialize.get_os_extension())
                    if not zr_bin:
                        logging.error("7zr not found, cannot extract archive")
                        return False
                    cmd = f'{zr_bin} x -aoa -o"{files_path}" -y "{file_localpath}"'
                    os.system(cmd)
        return True