From 8f2ef3fdaf4166d94fc7d8b3a3d2f5938dbec9dc Mon Sep 17 00:00:00 2001
From: mpgn <martial.puygrenier@gmail.com>
Date: Sat, 20 Jun 2020 13:20:27 -0400
Subject: [PATCH] Add color when smb status is not ACCESS_DENIED #391

---
 cme/logger.py        |  4 ++--
 cme/protocols/smb.py | 20 ++++++++++++++++----
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/cme/logger.py b/cme/logger.py
index 5efe044e..e72bc4b3 100755
--- a/cme/logger.py
+++ b/cme/logger.py
@@ -69,8 +69,8 @@ class CMEAdapter(logging.LoggerAdapter):
         msg, kwargs = self.process(u'{} {}'.format(colored("[*]", 'blue', attrs=['bold']), msg), kwargs)
         self.logger.info(msg, *args, **kwargs)
 
-    def error(self, msg, *args, **kwargs):
-        msg, kwargs = self.process(u'{} {}'.format(colored("[-]", 'red', attrs=['bold']), msg), kwargs)
+    def error(self, msg, color='red', *args, **kwargs):
+        msg, kwargs = self.process(u'{} {}'.format(colored("[-]", color, attrs=['bold']), msg), kwargs)
         self.logger.error(msg, *args, **kwargs)
 
     def debug(self, msg, *args, **kwargs):
diff --git a/cme/protocols/smb.py b/cme/protocols/smb.py
index b1f19ec0..e6142677 100755
--- a/cme/protocols/smb.py
+++ b/cme/protocols/smb.py
@@ -38,6 +38,17 @@ from traceback import format_exc
 smb_share_name = gen_random_string(5).upper()
 smb_server = None
 
+smb_error_status = [
+    "STATUS_ACCOUNT_DISABLED",
+    "STATUS_ACCOUNT_EXPIRED",
+    "STATUS_ACCOUNT_RESTRICTION",
+    "STATUS_INVALID_LOGON_HOURS",
+    "STATUS_INVALID_WORKSTATION",
+    "STATUS_LOGON_TYPE_NOT_GRANTED",
+    "STATUS_PASSWORD_EXPIRED",
+    "STATUS_PASSWORD_MUST_CHANGE"
+]
+
 def requires_smb_server(func):
     def _decorator(self, *args, **kwargs):
         global smb_server
@@ -90,7 +101,6 @@ def requires_smb_server(func):
 
     return wraps(func)(_decorator)
 
-
 class smb(connection):
 
     def __init__(self, args, db, host):
@@ -318,7 +328,8 @@ class smb(connection):
                                                         username,
                                                         password,
                                                         error,
-                                                        '({})'.format(desc) if self.args.verbose else ''))
+                                                        '({})'.format(desc) if self.args.verbose else ''),
+                                                        color='magenta' if error in smb_error_status else 'red')
 
             if error == 'STATUS_LOGON_FAILURE': self.inc_failed_login(username)
 
@@ -366,11 +377,12 @@ class smb(connection):
                 self.create_conn_obj()
         except SessionError as e:
             error, desc = e.getErrorString()
-            self.logger.error(u'{}\\{} {} {} {}'.format(domain,
+            self.logger.error(u'{}\\{}:{} {} {}'.format(domain,
                                                         username,
                                                         ntlm_hash,
                                                         error,
-                                                        '({})'.format(desc) if self.args.verbose else ''))
+                                                        '({})'.format(desc) if self.args.verbose else ''),
+                                                        color='magenta' if error in smb_error_status else 'red')
 
             if error == 'STATUS_LOGON_FAILURE': self.inc_failed_login(username)