TruffleHog import to MySQL

import.py

#!/usr/bin/env python3

"""
First intall the pip modules:
    pip3 install peewee
    pip3 install pymysql

Then run TruffleHog:
    docker run --rm -it -v "$PWD:/pwd" trufflesecurity/trufflehog:latest github file:///repo-dir --json > output.json
    docker run --rm -it -v "$PWD:/pwd" trufflesecurity/trufflehog:latest github --org=org-name --token=token --json > output.json

Then run this import script (check if DATABASe env is required)
    python 3 import.py output.json
"""

import json
import sys
import base64
import os
import datetime
from peewee import *
from playhouse.db_url import connect

db = connect(os.environ.get('DATABASE') or 'mysql://root:root@127.0.0.1:3306/trufflehog')

class BaseModel(Model):
    class Meta:
        database = db

class Finding(BaseModel):
    id = AutoField()
    detector = IntegerField() # https://github.com/trufflesecurity/trufflehog/issues/385
    repository = CharField()
    verified = BooleanField()
    commit = CharField()
    file = CharField()
    line = IntegerField()
    email = CharField()
    data = TextField()
    redacted = TextField()
    date = DateTimeField()
    created_at = DateTimeField(default=datetime.datetime.now)

db.create_tables([Finding])

if not sys.argv[1]:
    print('Usage: ' + sys.argv[0] + ' <truffleHog-output.json>')
    quit()

with open(sys.argv[1], 'r') as file:
    lines = file.readlines()
    for line in lines:

        item = json.loads(line)
        data = base64.b64decode(item['Raw'])
        git = item['SourceMetadata']['Data']['Git']

        finding, created = Finding.get_or_create(
            repository=git['repository'],
            commit=git['commit'],
            file=git['file'],
            line=git['line'],
        )

        print('[+] Finding: %s', item['Redacted'])

        if not created:
            print(' - Already in the DB')
            continue

        print(' - New finding')
        finding.detector = item['DetectorType']
        finding.verified = item['Verified']
        finding.email = git['email']
        finding.data = data
        finding.redacted = item['Redacted']
        finding.date = git['timestamp']
        finding.save()