reworked logger and added process logging

Pipfile

@@ -0,0 +1,12 @@
+[[source]]
+url = "https://pypi.org/simple"
+verify_ssl = true
+name = "pypi"
+
+[packages]
+psutil = "*"
+
+[dev-packages]
+
+[requires]
+python_version = "3.8"

activity_logger.py

@@ -9,16 +9,16 @@ def __init__(self, log_file_path):
       with open(self.log_file_path, 'w') as f:
         f.write('logs:')
 
-  def log_process_activity(self, **kwargs):
-    self.log_activity(activity_type='process', **kwargs)
+  # def log_process_activity(self, **kwargs):
+  #   self.log_activity(activity_type='process', **kwargs)
 
-  def log_file_activity(self, file_name, interaction, **kwargs):
-    full_file_path = os.path.realpath(file_name)
-    self.log_activity(
-      activity_type='file',
-      interaction=interaction,
-      full_file_path=full_file_path,
-      **kwargs)
+  # def log_file_activity(self, file_name, interaction, **kwargs):
+  #   full_file_path = os.path.realpath(file_name)
+  #   self.log_activity(
+  #     activity_type='file',
+  #     interaction=interaction,
+  #     full_file_path=full_file_path,
+  #     **kwargs)
 
   # def log_network_activity(self, dest_host, dest_port, src_host, src_port, bytes_tx, protocol, interaction, **kwargs):
   #   self.log_activity(
@@ -27,8 +27,18 @@ def log_file_activity(self, file_name, interaction, **kwargs):
   #     full_file_path=full_file_path,
   #     **kwargs)
 
-  def log_activity(self, **kwargs):
-    entries = {'timestamp': self.generate_timestamp(), **kwargs}
+  def log_activity(self, process, activity_type, timestamp=None, **kwargs):
+    if timestamp is None:
+      timestamp = ActivityLogger.generate_timestamp()
+    entries = {
+      'activity_type': activity_type,
+      'timestamp': timestamp,
+      'pid': process.pid,
+      'process_name': process.name(),
+      'process_cmd_line': process.cmdline(),
+      'initiated_by': process.username(),
+      **kwargs
+    }
     with open(self.log_file_path, 'a') as f:
       entries_formatted = '\n    '.join(['{}: "{}"'.format(k, v) for k, v in entries.items()])
       log_text = '\n  - {}'.format(entries_formatted)

main.py

@@ -1,60 +1,43 @@
 import argparse
 import os
 import time
-import subprocess
+import psutil
+from subprocess import DEVNULL
 from activity_logger import ActivityLogger
 
 parser = argparse.ArgumentParser(description='Trigger activity for EDR agent testing.')
-parser.add_argument('--file-path', metavar="/path/to/file.txt", required=True, type=str, dest='file_path')
-# parser.add_argument('--executable-path', metavar="/path/to/executable", required=True, type=str, dest='file_path')
+parser.add_argument('--test-file-path', metavar="/path/to/test_file.txt", required=True, type=str, dest='test_file_path')
+parser.add_argument('--log-file-path', metavar="/path/to/log.yaml", required=True, type=str, dest='log_file_path')
+parser.add_argument('--executable', metavar="\"/path/to/executable --with-flags\"", required=True, type=str, dest='executable')
 args = parser.parse_args()
 
-logger = ActivityLogger('logs/somelog.yaml')
+logger = ActivityLogger(args.log_file_path)
+current_process = psutil.Process(os.getpid())
+test_file_full_path = os.path.realpath(args.test_file_path)
 
 print('Creating file...')
-with open(args.file_path, 'x') as f:
-  logger.log_file_activity(
-    f.name,
-    'create',
-    initiated_by='???',
-    process_name='???',
-    process_cmd_line='???',
-    pid='???')
+with open(test_file_full_path, 'x') as f:
+  logger.log_activity(current_process, 'file', file_interaction='create', full_file_path=test_file_full_path)
 time.sleep(.35)
 print('Done.')
 
 print('Modifying file...')
-with open(args.file_path, 'w') as f:
+with open(test_file_full_path, 'w') as f:
   f.write("Modifying this file")
-  logger.log_file_activity(
-    f.name,
-    'modify',
-    initiated_by='???',
-    process_name='???',
-    process_cmd_line='???',
-    pid='???')
+  logger.log_activity(current_process, 'file', file_interaction='modify', full_file_path=test_file_full_path)
 time.sleep(.35)
 print('Done.')
 
 print('Deleting file...')
-os.remove(args.file_path)
-logger.log_file_activity(
-  f.name,
-  'delete',
-  initiated_by='???',
-  process_name='???',
-  process_cmd_line='???',
-  pid='???')
+os.remove(test_file_full_path)
+logger.log_activity(current_process, 'file', file_interaction='delete', full_file_path=test_file_full_path)
 time.sleep(.35)
 print('Done.')
 
 print('Starting process...')
-p = subprocess.Popen(['ls', '-al'])
-logger.log_process_activity(
-  initiated_by='???',
-  process_name='???',
-  process_cmd_line='???',
-  pid='???')
+process_args = args.executable.split()
+p = psutil.Popen(process_args, stdout=DEVNULL, stderr=DEVNULL)
+logger.log_activity(p, 'process', timestamp=p.create_time())
 time.sleep(.35)
 print('Done.')