Skip to content

Commit

Permalink
Moved LokiLogger class to separate file
Browse files Browse the repository at this point in the history
  • Loading branch information
Florian Roth committed Jun 12, 2017
1 parent 9844bb3 commit 9e72e02
Show file tree
Hide file tree
Showing 3 changed files with 178 additions and 172 deletions.
7 changes: 0 additions & 7 deletions lib/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -243,13 +243,6 @@ def removeNonAscii(string, stripit=False):

return nonascii


def getSyslogTimestamp():
date_obj = datetime.datetime.utcnow()
date_str = date_obj.strftime("%Y%m%dT%H:%M:%SZ")
return date_str


def getAge(filePath):
try:
stats=os.stat(filePath)
Expand Down
175 changes: 175 additions & 0 deletions lib/lokilogger.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,175 @@
# -*- coding: utf-8 -*-
#
# LOKI Logger

import sys, re
from colorama import Fore, Back, Style
from colorama import init
import codecs
import datetime

__version__ = '0.21.0'

# Logger Class -----------------------------------------------------------------
class LokiLogger():

no_log_file = False
log_file = "loki.log"
csv = False
hostname = "NOTSET"
alerts = 0
warnings = 0
notices = 0
only_relevant = False
debug = False

def __init__(self, no_log_file, log_file, hostname, csv, only_relevant, debug):
self.no_log_file = no_log_file
self.log_file = log_file
self.hostname = hostname
self.csv = csv
self.only_relevant = only_relevant
self.debug = debug

# Colorization ----------------------------------------------------
init()

# Welcome
if not self.csv:
self.print_welcome()

def log(self, mes_type, message):

# Remove all non-ASCII characters
# message = removeNonAsciiDrop(message)
codecs.register(lambda message: codecs.lookup('utf-8') if message == 'cp65001' else None)

if not self.debug and mes_type == "DEBUG":
return

# Counter
if mes_type == "ALERT":
self.alerts += 1
if mes_type == "WARNING":
self.warnings += 1
if mes_type == "NOTICE":
self.notices += 1

if self.only_relevant:
if mes_type not in ('ALERT', 'WARNING'):
return

# to stdout
self.log_to_stdout(message.encode('ascii', errors='replace'), mes_type)

# to file
if not self.no_log_file:
self.log_to_file(message, mes_type)

def log_to_stdout(self, message, mes_type):

# Prepare Message
#message = removeNonAsciiDrop(message)
codecs.register(lambda message: codecs.lookup('utf-8') if message == 'cp65001' else None)
message = message.encode(sys.stdout.encoding, errors='replace')

if self.csv:
print "{0},{1},{2},{3}".format(getSyslogTimestamp(),self.hostname,mes_type,message)

else:

try:

key_color = Fore.WHITE
base_color = Fore.WHITE+Back.BLACK
high_color = Fore.WHITE+Back.BLACK

if mes_type == "NOTICE":
base_color = Fore.CYAN+''+Back.BLACK
high_color = Fore.BLACK+''+Back.CYAN
elif mes_type == "INFO":
base_color = Fore.GREEN+''+Back.BLACK
high_color = Fore.BLACK+''+Back.GREEN
elif mes_type == "WARNING":
base_color = Fore.YELLOW+''+Back.BLACK
high_color = Fore.BLACK+''+Back.YELLOW
elif mes_type == "ALERT":
base_color = Fore.RED+''+Back.BLACK
high_color = Fore.BLACK+''+Back.RED
elif mes_type == "DEBUG":
base_color = Fore.WHITE+''+Back.BLACK
high_color = Fore.BLACK+''+Back.WHITE
elif mes_type == "ERROR":
base_color = Fore.MAGENTA+''+Back.BLACK
high_color = Fore.WHITE+''+Back.MAGENTA
elif mes_type == "RESULT":
if "clean" in message.lower():
high_color = Fore.BLACK+Back.GREEN
base_color = Fore.GREEN+Back.BLACK
elif "suspicious" in message.lower():
high_color = Fore.BLACK+Back.YELLOW
base_color = Fore.YELLOW+Back.BLACK
else:
high_color = Fore.BLACK+Back.RED
base_color = Fore.RED+Back.BLACK

# Colorize Type Word at the beginning of the line
type_colorer = re.compile(r'([A-Z]{3,})', re.VERBOSE)
mes_type = type_colorer.sub(high_color+r'[\1]'+base_color, mes_type)
# Break Line before REASONS
linebreaker = re.compile('(MD5:|SHA1:|SHA256:|MATCHES:|FILE:|FIRST_BYTES:|DESCRIPTION:|REASON_[0-9]+)', re.VERBOSE)
message = linebreaker.sub(r'\n\1', message)
# Colorize Key Words
colorer = re.compile('([A-Z_0-9]{2,}:)\s', re.VERBOSE)
message = colorer.sub(key_color+Style.BRIGHT+r'\1 '+base_color+Style.NORMAL, message)

# Print to console
if mes_type == "RESULT":
res_message = "\b\b%s %s" % (mes_type, message)
print base_color,res_message,Back.BLACK
print Fore.WHITE,Style.NORMAL
else:
sys.stdout.write("%s\b\b%s %s%s%s%s\n" % (base_color, mes_type, message, Back.BLACK,Fore.WHITE,Style.NORMAL))

except Exception, e:
traceback.print_exc()
print "Cannot print to cmd line - formatting error"

def log_to_file(self, message, mes_type):
try:
# Write to file
with codecs.open(self.log_file, "a", encoding='utf-8') as logfile:
if self.csv:
logfile.write(u"{0},{1},{2},{3}\n".format(getSyslogTimestamp(),self.hostname,mes_type,message))
else:
logfile.write(u"%s %s LOKI: %s: %s\n" % (getSyslogTimestamp(), self.hostname, mes_type.title(), message))
except Exception, e:
traceback.print_exc()
print "Cannot print to log file {0}".format(self.log_file)

def print_welcome(self):
print Back.GREEN + " ".ljust(79) + Back.BLACK + Fore.GREEN

print " __ ____ __ ______ "
print " / / / __ \/ //_/ _/ "
print " / /__/ /_/ / ,< _/ / "
print " /____/\____/_/|_/___/ "
print " ________ _____ ____ "
print " / _/ __ \/ ___/ / __/______ ____ ___ ___ ____ "
print " _/ // /_/ / /__ _\ \/ __/ _ `/ _ \/ _ \/ -_) __/ "
print " /___/\____/\___/ /___/\__/\_,_/_//_/_//_/\__/_/ "

print Fore.WHITE
print " (C) Florian Roth"
print " June 2017"
print " Version %s" % __version__
print " "
print " DISCLAIMER - USE AT YOUR OWN RISK"
print " "
print Back.GREEN + " ".ljust(79) + Back.BLACK
print Fore.WHITE+''+Back.BLACK

def getSyslogTimestamp():
date_obj = datetime.datetime.utcnow()
date_str = date_obj.strftime("%Y%m%dT%H:%M:%SZ")
return date_str
168 changes: 3 additions & 165 deletions loki.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,17 +32,15 @@
import re
import stat
import psutil
import codecs
from StringIO import StringIO
import signal as signal_module
import urllib2
import zipfile
import shutil
from StringIO import StringIO
from colorama import Fore, Back, Style
from colorama import init
from sys import platform as _platform

__version__ = '0.20.2'
# LOKI Modules
from lib.lokilogger import *

sys.stdout = codecs.getwriter('utf8')(sys.stdout)

Expand Down Expand Up @@ -1179,163 +1177,6 @@ def get_file_data(self, filePath):
return fileData


# Logger Class -----------------------------------------------------------------
class LokiLogger():

no_log_file = False
log_file = "loki.log"
csv = False
hostname = "NOTSET"
alerts = 0
warnings = 0
notices = 0
only_relevant = False
debug = False

def __init__(self, no_log_file, log_file, hostname, csv, only_relevant, debug):
self.no_log_file = no_log_file
self.log_file = log_file
self.hostname = hostname
self.csv = csv
self.only_relevant = only_relevant
self.debug = debug

# Welcome
if not self.csv:
self.print_welcome()

def log(self, mes_type, message):

# Remove all non-ASCII characters
# message = removeNonAsciiDrop(message)
codecs.register(lambda message: codecs.lookup('utf-8') if message == 'cp65001' else None)

if not args.debug and mes_type == "DEBUG":
return

# Counter
if mes_type == "ALERT":
self.alerts += 1
if mes_type == "WARNING":
self.warnings += 1
if mes_type == "NOTICE":
self.notices += 1

if self.only_relevant:
if mes_type not in ('ALERT', 'WARNING'):
return

# to stdout
self.log_to_stdout(message.encode('ascii', errors='replace'), mes_type)

# to file
if not self.no_log_file:
self.log_to_file(message, mes_type)

def log_to_stdout(self, message, mes_type):

# Prepare Message
#message = removeNonAsciiDrop(message)
codecs.register(lambda message: codecs.lookup('utf-8') if message == 'cp65001' else None)
message = message.encode(sys.stdout.encoding, errors='replace')

if self.csv:
print "{0},{1},{2},{3}".format(getSyslogTimestamp(),self.hostname,mes_type,message)

else:

try:

key_color = Fore.WHITE
base_color = Fore.WHITE+Back.BLACK
high_color = Fore.WHITE+Back.BLACK

if mes_type == "NOTICE":
base_color = Fore.CYAN+''+Back.BLACK
high_color = Fore.BLACK+''+Back.CYAN
elif mes_type == "INFO":
base_color = Fore.GREEN+''+Back.BLACK
high_color = Fore.BLACK+''+Back.GREEN
elif mes_type == "WARNING":
base_color = Fore.YELLOW+''+Back.BLACK
high_color = Fore.BLACK+''+Back.YELLOW
elif mes_type == "ALERT":
base_color = Fore.RED+''+Back.BLACK
high_color = Fore.BLACK+''+Back.RED
elif mes_type == "DEBUG":
base_color = Fore.WHITE+''+Back.BLACK
high_color = Fore.BLACK+''+Back.WHITE
elif mes_type == "ERROR":
base_color = Fore.MAGENTA+''+Back.BLACK
high_color = Fore.WHITE+''+Back.MAGENTA
elif mes_type == "RESULT":
if "clean" in message.lower():
high_color = Fore.BLACK+Back.GREEN
base_color = Fore.GREEN+Back.BLACK
elif "suspicious" in message.lower():
high_color = Fore.BLACK+Back.YELLOW
base_color = Fore.YELLOW+Back.BLACK
else:
high_color = Fore.BLACK+Back.RED
base_color = Fore.RED+Back.BLACK

# Colorize Type Word at the beginning of the line
type_colorer = re.compile(r'([A-Z]{3,})', re.VERBOSE)
mes_type = type_colorer.sub(high_color+r'[\1]'+base_color, mes_type)
# Break Line before REASONS
linebreaker = re.compile('(MD5:|SHA1:|SHA256:|MATCHES:|FILE:|FIRST_BYTES:|DESCRIPTION:|REASON_[0-9]+)', re.VERBOSE)
message = linebreaker.sub(r'\n\1', message)
# Colorize Key Words
colorer = re.compile('([A-Z_0-9]{2,}:)\s', re.VERBOSE)
message = colorer.sub(key_color+Style.BRIGHT+r'\1 '+base_color+Style.NORMAL, message)

# Print to console
if mes_type == "RESULT":
res_message = "\b\b%s %s" % (mes_type, message)
print base_color,res_message,Back.BLACK
print Fore.WHITE,Style.NORMAL
else:
sys.stdout.write("%s\b\b%s %s%s%s%s\n" % (base_color, mes_type, message, Back.BLACK,Fore.WHITE,Style.NORMAL))

except Exception, e:
traceback.print_exc()
print "Cannot print to cmd line - formatting error"

def log_to_file(self, message, mes_type):
try:
# Write to file
with codecs.open(self.log_file, "a", encoding='utf-8') as logfile:
if self.csv:
logfile.write(u"{0},{1},{2},{3}\n".format(getSyslogTimestamp(),self.hostname,mes_type,message))
else:
logfile.write(u"%s %s LOKI: %s: %s\n" % (getSyslogTimestamp(), self.hostname, mes_type.title(), message))
except Exception, e:
traceback.print_exc()
print "Cannot print to log file {0}".format(self.log_file)

def print_welcome(self):
print Back.GREEN + " ".ljust(79) + Back.BLACK + Fore.GREEN

print " __ ____ __ ______ "
print " / / / __ \/ //_/ _/ "
print " / /__/ /_/ / ,< _/ / "
print " /____/\____/_/|_/___/ "
print " ________ _____ ____ "
print " / _/ __ \/ ___/ / __/______ ____ ___ ___ ____ "
print " _/ // /_/ / /__ _\ \/ __/ _ `/ _ \/ _ \/ -_) __/ "
print " /___/\____/\___/ /___/\__/\_,_/_//_/_//_/\__/_/ "

print Fore.WHITE
print " (C) Florian Roth"
print " April 2017"
print " Version %s" % __version__
print " "
print " DISCLAIMER - USE AT YOUR OWN RISK"
print " "
print Back.GREEN + " ".ljust(79) + Back.BLACK
print Fore.WHITE+''+Back.BLACK


def walk_error(err):
if "Error 3" in str(err):
logger.log("ERROR", str(err))
Expand Down Expand Up @@ -1466,9 +1307,6 @@ def signal_handler(signal_name, frame):

args = parser.parse_args()

# Colorization ----------------------------------------------------
init()

# Remove old log file
if os.path.exists(args.l):
os.remove(args.l)
Expand Down

0 comments on commit 9e72e02

Please sign in to comment.