Merge pull request #6 from virtru/2.x-RC

v2.0.0

CHANGELOG.md

@@ -6,6 +6,11 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 
 ## [Unreleased](https://github.com/virtru/virtru-audit-export-client/compare/master...HEAD)
 
+## [2.0.0](https://github.com/virtru/audit-export-client/pull/5) - 2019-03-25
+## Added
+ - Changed cli logic for using elastic search api
+ - Using cursor logic instead bookmark
+
 ## [1.1.0](https://github.com/virtru/audit-export-client/pull/4) - 2018-10-31
 ## Added
  - Added tests

README.md

@@ -48,15 +48,16 @@ to export audit records to csv, provide a path to a folder
 to export audit records to syslog, provide a host and a port
 ### `--sysloghost=0.0.0.0 --syslogport=514`
 
-to pull records since the last time the script was run, set the bookmark option
-### `--bookmark` or `-b`
-you can also provide a `bookmark.ini` file, in  the `.auditexport` directory, which will tell the script where to start pulling records
+to pull records since the last time the script was run, set the cursor option
+### `--cursor` or `-c`
+you can also provide a `cursor.ini` file, in  the `.auditexport` directory, which will tell the script where to start pulling records
 
 ```ini
-#bookmark.ini
+#cursor.ini
 
 [next-page-start-key]
-nextpagestartkey=<next-page-start-index>
+nextpagecursor=<next-page-cursor>
+lastrecordsaved=<recordId>
 
 ```
 

auditexport/auditclient/auditclient.py

@@ -16,6 +16,8 @@
 VJWT_TTL_SECONDS = 300.0
 API_HOST = 'audit.virtru.com'
 API_PATH = '/api/messages'
+CLIENT_NAME = 'AuditPythonClient:v2.0.0'
+API_VERSION = 'v2'
 
 
 class AuditClient:
@@ -64,7 +66,9 @@ def fetchRecords(self, req):
 
         headers = {
             'Authorization': 'VJWTv1.0.0 ' + vjwtString.decode(),
-            'Connection': 'close'
+            'Connection': 'close',
+            'X-Virtru-Client': CLIENT_NAME,
+            'Accept-Version': API_VERSION
         }
         apiUrl = self.apiHost + self.apiPath
 

auditexport/cli.py

@@ -46,12 +46,17 @@ def main():
                         dest='syslogport',
                         default='514',
                         required=False)
-    parser.add_argument('-b', '--bookmark',
-                        help='Start from last bookmark',
-                        dest='useBookMark',
+    parser.add_argument('-c', '--cursor',
+                        help='Start from last cursor',
+                        dest='useCursor',
                         default=False,
                         required=False,
                         action='store_true')
+    parser.add_argument('-l', '--limit',
+                        help='Number of records we pull for each chunk when we use --cursor. Default is 100',
+                        dest='limit',
+                        default=100,
+                        required=False)
     parser.add_argument('-v', '--verbose',
                         help='Verbose option',
                         dest='verbose',
@@ -87,12 +92,13 @@ def main():
 
 def process(args, auditclient, utils):
 
-    logger.debug('fetching bookmark.......')
+    logger.debug('fetching cursor.......')
 
-    bookMark = utils.getNextPageStartKey()
-    nextPageStartKey = None if not bookMark else bookMark['nextpagestartkey']
+    cursor = utils.getnextPageCursor()
+    nextPageCursor = None if not cursor else cursor['nextPageCursor']
+    lastRecordId = None if not cursor else cursor['lastRecordSaved']
 
-    logger.debug('nextpagestartkey: %s' % (nextPageStartKey))
+    logger.debug('nextPageCursor: %s' % (nextPageCursor))
 
     queryStart = args.startDate
     queryEnd = args.endDate
@@ -105,9 +111,10 @@ def process(args, auditclient, utils):
     csvFolderPath = args.csv
     syslogHost = args.sysloghost
     syslogPort = args.syslogport
-    useBookMark = args.useBookMark
+    useCursor = args.useCursor
+    limit = args.limit
 
-    logger.debug('usebookmark: %s' % (useBookMark))
+    logger.debug('useCursor: %s' % (useCursor))
 
     # Syslog logger
     syslogger = None if syslogHost is None else utils.configSysLogger(
@@ -117,40 +124,51 @@ def process(args, auditclient, utils):
         'method': 'GET',
         'query': {
             'start': queryStart,
-            'end': queryEnd
+            'end': queryEnd,
+            'sort': 'timestamp:asc',
         }
     }
 
-    if(nextPageStartKey and useBookMark):
-        req['query']['nextPageStartKey'] = nextPageStartKey
+    if(nextPageCursor and useCursor):
+        req['query']['cursor'] = nextPageCursor
+
+    req['query']['limit'] = limit
 
     hasMore = True
     iteration = 1
 
     logger.debug('fetching audit records....')
     while hasMore:
-        records = auditclient.fetchRecords(req)
-        if(jsonFolderPath and records['docs']):
-            utils.exportToJson(jsonFolderPath, records['docs'])
-        if(csvFolderPath and records['docs']):
-            utils.exportToCsv(csvFolderPath, records['docs'])
-        if(syslogHost is not None and records['docs']):
+        payload = auditclient.fetchRecords(req)
+        records = payload['data'] if not cursor else utils.checkRecords(
+            payload['data'], lastRecordId)
+        if(len(records)):
+            lastRecordId = records[-1]['recordId']
+        else:
+            hasMore = False
+            break
+
+        if(jsonFolderPath and len(records)):
+            utils.exportToJson(jsonFolderPath, records)
+        if(csvFolderPath and len(records)):
+            utils.exportToCsv(csvFolderPath, records)
+        if(syslogHost is not None and len(records)):
             utils.exportToSysLog(syslogHost, syslogPort,
-                                 syslogger, records['docs'])
+                                 syslogger, records)
 
-        if 'nextPageStartKey' in records:
-            logger.debug('found nextpagestartkey')
-            nextPageStartKey = records['nextPageStartKey']
-            req['query']['nextPageStartKey'] = nextPageStartKey
+        if 'after' in payload['cursor']:
+            logger.debug('found next cursor')
+            nextPageCursor = payload['cursor']['after']
+            req['query']['cursor'] = nextPageCursor
         else:
             hasMore = False
-            if records['docs']:
-                nextPageStartKey = records['docs'][-1]['recordId']
 
-        if(useBookMark):
-            utils.saveNextPageStartKey(nextPageStartKey)
+        if(useCursor):
+            utils.saveNextPageCursor(nextPageCursor, lastRecordId)
+
+        cursorToPrint = str(nextPageCursor) if hasMore else 'None'
 
         print('Iteration :' + str(iteration) + '\t\t' + 'Items: ' +
-              str(len(records['docs'])) + '\t\t' + 'NextPageStartKey: ' + str(nextPageStartKey))
+              str(len(records)) + '\t\t' + 'nextPageCursor: ' + cursorToPrint)
         iteration += 1
     print('All records exported!!!!')

auditexport/utils.py

@@ -13,7 +13,7 @@
 from .auditclient.errors import AuditClientError
 
 EXPORT_DIR = '.auditexport'
-BOOK_MARK_PATH = '%s/bookmark.ini' % (EXPORT_DIR)
+CURSOR_PATH = '%s/cursor.ini' % (EXPORT_DIR)
 
 
 logger = logging.getLogger(__name__)
@@ -72,52 +72,69 @@ def getConfig(configFile=''):
         'apiPath': apiPath
     }
 
+def checkRecords(records=[], savedId=None):
+    if not savedId:
+        return records
+    for i in range(len(records)):
+        if records[i]['recordId'] == savedId:
+            return records[i+1:]
+    return records
 
-def getNextPageStartKey():
-    bookmark = configparser.ConfigParser()
-    bookmark.read(BOOK_MARK_PATH)
+
+def getnextPageCursor():
+    cursor = configparser.ConfigParser()
+    cursor.read(CURSOR_PATH)
 
     # Config Parser returns an empty dataset if file does not exist
-    if len(bookmark) <= 1:
+    if len(cursor) <= 1:
         return None
     else:
-        return bookmark['next-page-start-key']
+        return cursor['next-page-cursor']
 
+def flattenObjectProp(record):
+    editedRecord = {**record, **record['object']}
+    del(editedRecord['object'])
+    return editedRecord
 
-def saveNextPageStartKey(nextPageStartKey):
-    logger.debug('saving nexpagestartkey.....')
+def saveNextPageCursor(nextPageCursor, lastRecordSaved):
+    logger.debug('saving next-page-cursor.....')
 
-    bookMarkConfig = configparser.ConfigParser()
-    bookMarkConfig['next-page-start-key'] = {
-        'nextPageStartKey': nextPageStartKey}
-    os.makedirs(os.path.dirname(BOOK_MARK_PATH), exist_ok=True)
-    with open(BOOK_MARK_PATH, 'w') as bookMarkFile:
-        bookMarkConfig.write(bookMarkFile)
+    cursorConfig = configparser.ConfigParser()
+    cursorConfig['next-page-cursor'] = {
+        'nextPageCursor': nextPageCursor or '',
+        'lastRecordSaved': lastRecordSaved
+    }
+    os.makedirs(os.path.dirname(CURSOR_PATH), exist_ok=True)
+    with open(CURSOR_PATH, 'w') as cursorFile:
+        cursorConfig.write(cursorFile)
 
 
 def exportToJson(pathToFolder, records):
     logger.debug('exporting records to json.....')
+    preparedRecords = list(map(lambda record: flattenObjectProp(record), records))
 
     fileName = str(datetime.datetime.utcnow().isoformat()) + ".json"
     fn = os.path.join(pathToFolder, fileName)
     with open(fn, "w") as f:
-        json.dump(records, f, sort_keys=True,
+        json.dump(preparedRecords, f, sort_keys=True,
                   indent=4, separators=(',', ': '))
 
 
 def exportToCsv(pathToFolder, records):
     logger.debug('exporting records to csv.....')
+    preparedRecords = list(map(lambda record: flattenObjectProp(record), records))
 
-    for record in records:
+    for record in preparedRecords:
         auditType = record['type']
         fileName = auditType + ".csv"
         __writeCsvFile(auditType, pathToFolder, fileName, record)
 
 
 def exportToSysLog(host, port, syslogger, records):
     logger.debug('exporting to records to syslog......')
+    preparedRecords = list(map(lambda record: flattenObjectProp(record), records))
 
-    for record in records:
+    for record in preparedRecords:
         # Flatten out dictionary
         formattedRecord = __flatten(record)