Skip to content

Commit

Permalink
Merge pull request #142 from thorinaboenke/132_increasing_logs
Browse files Browse the repository at this point in the history
132 append to log files
whotwagner authored Dec 6, 2024
2 parents 5213735 + c74df88 commit ea2ee1b
Showing 4 changed files with 60 additions and 13 deletions.
3 changes: 2 additions & 1 deletion docs/source/basic.rst
Original file line number Diff line number Diff line change
@@ -7,7 +7,7 @@ AttackMate ships with a executable stub called "attackmate" that can be called l
::

attackmate -h
usage: attackmate [-h] --config CONFIG [--debug] [--version]
usage: attackmate [-h] --config CONFIG [--debug] [--version] [--json] [--append_logs]

AttackMate is an attack orchestration tool that executes full attack-chains based on playbooks.

@@ -17,6 +17,7 @@ AttackMate ships with a executable stub called "attackmate" that can be called l
--debug Enable verbose output
--version show program's version number and exit
--json log commands to attackmate.json
--append_logs append logs to attackmate.log, output.log and attackmate.json instead of overwriting

(Austrian Institute of Technology) https://aecid.ait.ac.at Version: 0.2.0

12 changes: 9 additions & 3 deletions src/attackmate/__main__.py
Original file line number Diff line number Diff line change
@@ -20,16 +20,22 @@ def parse_args():
parser.add_argument('--config', help='Configfile in yaml-format')
parser.add_argument('--debug', action='store_true', default=False, help='Enable verbose output')
parser.add_argument('--json', action='store_true', default=False, help='log commands to attackmate.json')
parser.add_argument(
'--append_logs',
action='store_true',
default=False,
help='append logs to attackmate.log, output.log and attackmate.json instead of overwriting',
)
parser.add_argument('--version', action='version', version=__version_string__)
parser.add_argument('playbook', help='Playbook in yaml-format')
return parser.parse_args()


def main():
args = parse_args()
logger = initialize_logger(args.debug)
initialize_output_logger(args.debug)
initialize_json_logger(args.json)
logger = initialize_logger(args.debug, args.append_logs)
initialize_output_logger(args.debug, args.append_logs)
initialize_json_logger(args.json, args.append_logs)
hacky = AttackMate(parse_playbook(args.playbook, logger), parse_config(args.config, logger))
sys.exit(hacky.main())

25 changes: 16 additions & 9 deletions src/attackmate/logging_setup.py
Original file line number Diff line number Diff line change
@@ -4,21 +4,30 @@
from colorlog import ColoredFormatter


def initialize_output_logger(debug: bool):
def create_file_handler(
file_name: str, append_logs: bool, formatter: logging.Formatter
) -> logging.FileHandler:
mode = 'a' if append_logs else 'w'
file_handler = logging.FileHandler(file_name, mode=mode)
file_handler.setFormatter(formatter)

return file_handler


def initialize_output_logger(debug: bool, append_logs: bool):
output_logger = logging.getLogger('output')
if debug:
output_logger.setLevel(logging.DEBUG)
else:
output_logger.setLevel(logging.INFO)
file_handler = logging.FileHandler('output.log', mode='w')
formatter = logging.Formatter(
'--- %(asctime)s %(levelname)s: ---\n\n%(message)s', datefmt='%Y-%m-%d %H:%M:%S'
)
file_handler.setFormatter(formatter)
file_handler = create_file_handler('output.log', append_logs, formatter)
output_logger.addHandler(file_handler)


def initialize_logger(debug: bool):
def initialize_logger(debug: bool, append_logs: bool):
playbook_logger = logging.getLogger('playbook')
if debug:
playbook_logger.setLevel(logging.DEBUG)
@@ -33,22 +42,20 @@ def initialize_logger(debug: bool):

# plain text output
playbook_logger.addHandler(console_handler)
file_handler = logging.FileHandler('attackmate.log', mode='w')
formatter = logging.Formatter('%(asctime)s %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
file_handler.setFormatter(formatter)
file_handler = create_file_handler('attackmate.log', append_logs, formatter)
playbook_logger.addHandler(file_handler)

return playbook_logger


def initialize_json_logger(json: bool):
def initialize_json_logger(json: bool, append_logs: bool):
if not json:
return None
json_logger = logging.getLogger('json')
json_logger.setLevel(logging.DEBUG)
file_handler = logging.FileHandler('attackmate.json', mode='w')
formatter = logging.Formatter('%(message)s')
file_handler.setFormatter(formatter)
file_handler = create_file_handler('attackmate.json', append_logs, formatter)
json_logger.addHandler(file_handler)

return json_logger
33 changes: 33 additions & 0 deletions test/units/test_logging.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
import pytest
from unittest.mock import patch, MagicMock
from attackmate.logging_setup import create_file_handler


@patch('attackmate.logging_setup.logging.FileHandler')
def test_create_file_handler_append_mode(MockFileHandler):
mock_handler = MagicMock()
MockFileHandler.return_value = mock_handler

file_name = 'test.log'
append_logs = True
formatter = MagicMock()

create_file_handler(file_name, append_logs, formatter)

MockFileHandler.assert_called_with(file_name, mode='a')
mock_handler.setFormatter.assert_called_with(formatter)


@patch('attackmate.logging_setup.logging.FileHandler')
def test_create_file_handler_write_mode(MockFileHandler):
mock_handler = MagicMock()
MockFileHandler.return_value = mock_handler

file_name = 'test.log'
append_logs = False
formatter = MagicMock()

create_file_handler(file_name, append_logs, formatter)

MockFileHandler.assert_called_with(file_name, mode='w')
mock_handler.setFormatter.assert_called_with(formatter)

0 comments on commit ea2ee1b

Please sign in to comment.