From 42ccf773b6bfcd4841ab6081eaea6801fd566635 Mon Sep 17 00:00:00 2001
From: Helvio Junior <helvio_junior@hotmail.com>
Date: Tue, 30 Jul 2024 09:59:03 -0300
Subject: [PATCH] Improve output

---
 knowsmore/cmd/credentials.py | 53 +++++++++++++++++++++++++++---------
 knowsmore/util/tools.py      | 10 ++++++-
 2 files changed, 49 insertions(+), 14 deletions(-)

diff --git a/knowsmore/cmd/credentials.py b/knowsmore/cmd/credentials.py
index 35bcb05..02b2d63 100644
--- a/knowsmore/cmd/credentials.py
+++ b/knowsmore/cmd/credentials.py
@@ -23,6 +23,7 @@ class Credentials(CmdBase):
     db = None
     out_file = None
     out_path = None
+    out_file_json = None
 
     def __init__(self):
         super().__init__('credentials', 'Show cracked credentials')
@@ -32,6 +33,13 @@ def add_flags(self, flags: _ArgumentGroup):
                            action='store',
                            default='',
                            dest=f'out_file',
+                           help=Color.s(
+                               'Output file to save TXT file'))
+
+        flags.add_argument('--save-to-json',
+                           action='store',
+                           default='',
+                           dest=f'out_file_json',
                            help=Color.s(
                                'Output file to save JSON data'))
 
@@ -56,6 +64,15 @@ def load_from_arguments(self, args: Namespace) -> bool:
                     self.out_file))
                 exit(1)
 
+        if args.out_file_json is not None and args.out_file_json.strip() != '':
+            self.out_file_json = Path(args.out_file_json).absolute()
+
+        if self.out_file is not None:
+            if os.path.exists(self.out_file):
+                Logger.pl('{!} {R}error: out file ({O}%s{R}) already exists {W}\r\n' % (
+                    self.out_file))
+                exit(1)
+
         if args.out_path is not None and args.out_path.strip() != '':
             self.out_path = Path(args.out_path).absolute()
 
@@ -95,10 +112,10 @@ def run(self):
                     'rows': rows
                 })
 
-        if self.out_file is not None:
-            Color.pl('{?} {W}{D}Credentials saved at {W}{C}%s{W}{D}{W}' % self.out_file)
+        if self.out_file_json is not None:
+            Color.pl('{?} {W}{D}Credentials saved at {W}{C}%s{W}{D}{W}' % self.out_file_json)
 
-            with open(self.out_file, "a", encoding="UTF-8") as text_file:
+            with open(self.out_file_json, "a", encoding="UTF-8") as text_file:
                 text_file.write(json.dumps(
                     {
                         'data': data,
@@ -110,7 +127,7 @@ def run(self):
                     }
                 ))
 
-        elif self.out_path is not None:
+        elif self.out_path is not None or self.out_file is not None:
 
             for i, d in enumerate(data):
                 name = f"{i:03}_{Tools.sanitize_filename(d['description'])}"
@@ -129,15 +146,25 @@ def run(self):
                 else:
                     file_data += Tools.get_ansi_tabulated(d['rows'])
 
-                o = Ansi2Image(0, 0, font_name=Ansi2Image.get_default_font_name(), font_size=13)
-                o.loads(file_data)
-                o.min_margin = 10
-                o.max_margin = 30
-                o.calc_size(margin=0.01)
-                o.save_image(os.path.join(self.out_path, f'{name}.png'), format='PNG')
-
-                #with open(os.path.join(self.out_path, f'{name}.ansi.txt'), 'wb') as f:
-                #    f.write(file_data.encode('utf-8', 'ignore'))
+                if self.out_path is not None:
+                    o = Ansi2Image(0, 0, font_name=Ansi2Image.get_default_font_name(), font_size=13)
+                    o.loads(file_data)
+                    o.min_margin = 10
+                    o.max_margin = 30
+                    o.calc_size(margin=0.01)
+                    o.save_image(os.path.join(self.out_path, f'{name}.png'), format='PNG')
+
+                if self.out_file is not None:
+                    name = str(self.out_file).replace(Path(self.out_file).suffix, "").rstrip(". ")
+
+                    with open(f'{name}.ansi.txt', 'ab') as f:
+                        f.write(file_data.encode('utf-8', 'ignore'))
+                        f.write(b"\n\n")
+
+                    with open(f'{name}.txt', 'ab') as f:
+                        f.write(f"{d['description']}\n".encode('utf-8', 'ignore'))
+                        f.write(Tools.get_tabulated(d['rows']).encode('utf-8', 'ignore'))
+                        f.write(b"\n\n")
 
         else:
 
diff --git a/knowsmore/util/tools.py b/knowsmore/util/tools.py
index 331a1a7..5b3c5c6 100644
--- a/knowsmore/util/tools.py
+++ b/knowsmore/util/tools.py
@@ -198,4 +198,12 @@ def json_serial(obj):
         if isinstance(obj, bytes):
             return base64.b64encode(obj).decode("UTF-8")
 
-        raise TypeError("Type %s not serializable" % type(obj))
\ No newline at end of file
+        raise TypeError("Type %s not serializable" % type(obj))
+
+    @staticmethod
+    def escape_ansi(text):
+        if text is None:
+            return ''
+
+        pattern = re.compile(r'(\x9B|\x1B\[)[0-?]*[ -/]*[@-~]')
+        return pattern.sub('', text)