# Impacket - Collection of Python classes for working with network protocols.
#
# Copyright Fortra, LLC and its affiliated companies 
#
# All rights reserved.
#
# This software is provided under a slightly modified version
# of the Apache Software License. See the accompanying LICENSE file
# for more information.
#
# Description:
#   This logger is intended to be used by impacket instead
#   of printing directly. This will allow other libraries to use their
#   custom logging implementation.
#

import logging
import sys
from impacket import version
from impacket.examples.ntlmrelayx.utils.identity_log import IdentityFilter

# This module can be used by scripts using the Impacket library 
# in order to configure the root logger to output events 
# generated by the library with a predefined format

# If the scripts want to generate log entries, they can write
# directly to the root logger (logging.info, debug, etc).

class ImpacketFormatter(logging.Formatter):
  '''
  Prefixing logged messages through the custom attribute 'bullet'.
  '''
  def __init__(self):
      logging.Formatter.__init__(self,'%(bullet)s %(identity)s%(message)s', None)

  def format(self, record):
    if record.levelno == logging.INFO:
      record.bullet = '[*]'
    elif record.levelno == logging.DEBUG:
      record.bullet = '[+]'
    elif record.levelno == logging.WARNING:
      record.bullet = '[!]'
    else:
      record.bullet = '[-]'

    if not hasattr(record, 'identity'): record.identity = ''

    return logging.Formatter.format(self, record)

class ImpacketFormatterTimeStamp(ImpacketFormatter):
  '''
  Prefixing logged messages through the custom attribute 'bullet'.
  '''
  def __init__(self):
      logging.Formatter.__init__(self,'[%(asctime)-15s] %(bullet)s %(identity)s%(message)s', None)

  def formatTime(self, record, datefmt=None):
      return ImpacketFormatter.formatTime(self, record, datefmt="%Y-%m-%d %H:%M:%S")

def init(ts=False, debug=False):
    # We add a StreamHandler and formatter to the root logger
    handler = logging.StreamHandler(sys.stdout)

    if ts:
        handler.setFormatter(ImpacketFormatterTimeStamp())
    else:
        handler.setFormatter(ImpacketFormatter())

    handler.addFilter(IdentityFilter())

    logging.getLogger().addHandler(handler)

    if debug is True:
        logging.getLogger().setLevel(logging.DEBUG)
        # Print the Library's installation path
        logging.debug(version.getInstallationPath())
    else:
        logging.getLogger().setLevel(logging.INFO)
        logging.getLogger('impacket.smbserver').setLevel(logging.ERROR)