Skip to content

Commit

Permalink
pydriller
Browse files Browse the repository at this point in the history
  • Loading branch information
babenek committed Sep 28, 2024
1 parent 754f893 commit 0894855
Show file tree
Hide file tree
Showing 8 changed files with 268 additions and 58 deletions.
3 changes: 3 additions & 0 deletions .mypy.ini
Original file line number Diff line number Diff line change
Expand Up @@ -49,5 +49,8 @@ ignore_missing_imports = True
[mypy-docx.*]
ignore_missing_imports = True

[mypy-pydriller.*]
ignore_missing_imports = True

[mypy-base62.*]
ignore_missing_imports = True
102 changes: 100 additions & 2 deletions credsweeper/__main__.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,15 @@
import base64
import binascii
import hashlib
import io
import logging
import os
import sys
import time
from argparse import ArgumentParser, ArgumentTypeError, Namespace
from typing import Any, Union, Optional, Dict
from typing import Any, Union, Optional, Dict, List, Tuple

from pydriller import Repository

from credsweeper import __version__
from credsweeper.app import APP_PATH, CredSweeper
Expand Down Expand Up @@ -116,6 +121,17 @@ def get_arguments() -> Namespace:
const="log.yaml",
dest="export_log_config",
metavar="PATH")
group.add_argument("--git", nargs="+", help="git repo to scan", dest="git", metavar="PATH")
parser.add_argument("--commits",
help="scan git repo for N commits only",
type=positive_int,
dest="commits",
default=0,
metavar="POSITIVE_INT")
parser.add_argument("--branch",
help="scan git repo for single branch, otherwise - all branches were scanned (slow)",
dest="branch",
type=str)
parser.add_argument("--rules",
help="path of rule config file (default: credsweeper/rules/config.yaml). "
f"severity:{[i.value for i in Severity]} "
Expand Down Expand Up @@ -320,9 +336,84 @@ def scan(args: Namespace, content_provider: AbstractProvider, json_filename: Opt
return credsweeper.run(content_provider=content_provider)
except Exception as exc:
logger.critical(exc, exc_info=True)
logger.exception(exc)
return -1


def scan_git(args: Namespace) -> Tuple[int, int, int]:
"""Scan repository for branches and commits
Returns:
total credentials found
total scanned branches
total scanned commits
"""
total_credentials = 0
total_branches = 0
total_commits = 0
try:
sha1git = hashlib.sha1(str(args.git).encode()).digest()
repo_hash = base64.b32encode(sha1git).decode("ascii")
journal_filename = f"{repo_hash}.json"
logger.info(f"{args.git} sha1 in base32 {repo_hash}")
repo_journal = Util.json_load(journal_filename)
if not isinstance(repo_journal, dict):
with open(journal_filename, "w") as f:
f.write("{}")
repo_journal = {"repo": args.git}
credsweeper = CredSweeper(rule_path=args.rule_path,
config_path=args.config_path,
api_validation=args.api_validation,
sort_output=args.sort_output,
use_filters=args.no_filters,
pool_count=args.jobs,
ml_batch_size=args.ml_batch_size,
ml_threshold=args.ml_threshold,
ml_providers=args.ml_providers,
find_by_ext=args.find_by_ext,
depth=args.depth,
doc=args.doc,
severity=args.severity,
size_limit=args.size_limit,
log_level=args.log)
repository = Repository(args.git, only_in_branch=args.branch)
for commit in repository.traverse_commits():
if commit.hash in repo_journal:
logger.debug(f"Skip already scanned commit: {commit.hash}")
continue
logger.info(f"Scan commit: {commit.hash}")
paths: List[Tuple[str, io.BytesIO]] = []
for file in commit.modified_files:
logger.info(f"FILE: {file.old_path} -> {file.new_path}")
try:
if file.new_path is not None:
_io = io.BytesIO(file.content)
paths.append((file.filename, _io))
except ValueError as exc:
logger.error("Possible missed submodule:%s", str(exc))
provider = FilesProvider(paths)
if args.json_filename:
ext = Util.get_extension(args.json_filename, False)
credsweeper.json_filename = f"{args.json_filename[:-len(ext)]}.{commit.hash}{ext}"
if args.xlsx_filename:
ext = Util.get_extension(args.xlsx_filename, False)
credsweeper.xlsx_filename = f"{args.xlsx_filename[:-len(ext)]}.{commit.hash}{ext}"

commit_cred_number = credsweeper.run(provider)
if credsweeper.is_ml_validator_inited:
# reset not-pickled object for multiprocess
credsweeper.ml_validator = None
credsweeper.credential_manager.candidates.clear()
total_credentials += commit_cred_number
total_commits += 1
repo_journal[commit.hash] = commit_cred_number
Util.json_dump(repo_journal, journal_filename)
total_branches += 1
except Exception as exc:
logger.critical(exc, exc_info=True)
return -1, total_branches, total_commits
return total_credentials, total_branches, total_commits


def main() -> int:
"""Main function"""
result = EXIT_FAILURE
Expand All @@ -331,7 +422,7 @@ def main() -> int:
if args.banner:
print(f"CredSweeper {__version__} crc32:{check_integrity():08x}")
Logger.init_logging(args.log, args.log_config_path)
logger.info(f"Init CredSweeper object with arguments: {args}")
logger.info(f"Init CredSweeper object with arguments: {args} CWD: {os.getcwd()}")
summary: Dict[str, int] = {}
if args.path:
logger.info(f"Run analyzer on path: {args.path}")
Expand All @@ -354,6 +445,13 @@ def main() -> int:
summary["Deleted File Credentials"] = del_credentials_number
if 0 <= add_credentials_number and 0 <= del_credentials_number:
result = EXIT_SUCCESS
elif args.git:
logger.info(f"Run analyzer on GIT: {args.git}")
credentials_number, branches_number, commits_number = scan_git(args)
summary[
f"Detected Credentials in {branches_number} branches and {commits_number} commits "] = credentials_number
if 0 <= credentials_number:
result = EXIT_SUCCESS
elif args.export_config:
logging.info(f"Exporting default config to file: {args.export_config}")
config_dict = Util.json_load(APP_PATH / "secret" / "config.json")
Expand Down
18 changes: 16 additions & 2 deletions credsweeper/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ def __init__(self,
sort_output: bool = False,
use_filters: bool = True,
pool_count: int = 1,
ml_batch_size: Optional[int] = None,
ml_batch_size: Optional[int] = 16,
ml_threshold: Union[float, ThresholdPreset] = ThresholdPreset.medium,
ml_config: Union[None, str, Path] = None,
ml_model: Union[None, str, Path] = None,
Expand Down Expand Up @@ -187,11 +187,18 @@ def _use_ml_validation(self) -> bool:

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

@property
def is_ml_validator_inited(self) -> bool:
"""method to check whether ml_validator was inited without creation"""
return bool(self.__ml_validator)

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

@property
def ml_validator(self) -> MlValidator:
"""ml_validator getter"""
from credsweeper.ml_model import MlValidator
if not self.__ml_validator:
if not self.is_ml_validator_inited:
self.__ml_validator: MlValidator = MlValidator(
threshold=self.ml_threshold, #
ml_config=self.ml_config, #
Expand Down Expand Up @@ -381,6 +388,7 @@ def post_processing(self) -> None:
if ml_cred_groups:
logger.info(f"Run ML Validation for {len(ml_cred_groups)} groups")
is_cred, probability = self.ml_validator.validate_groups(ml_cred_groups, self.ml_batch_size)
logger.info(f"DONE ML Validation for {len(is_cred)} results")
for i, (_, group_candidates) in enumerate(ml_cred_groups):
for candidate in group_candidates:
if candidate.use_ml:
Expand All @@ -404,6 +412,12 @@ def export_results(self) -> None:

credentials = self.credential_manager.get_credentials()

if credentials:
logger.info(f"Exporting {len(credentials)} credentials")
else:
logger.info("No credentials were found")
return

if self.sort_output:
credentials.sort(key=lambda x: ( #
x.line_data_list[0].path, #
Expand Down
4 changes: 4 additions & 0 deletions credsweeper/ml_model/ml_validator.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from credsweeper.credentials import Candidate, CandidateKey
import credsweeper.ml_model.features as features
from credsweeper.utils import Util
import psutil

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -238,6 +239,9 @@ def validate_groups(self, group_list: List[Tuple[CandidateKey, List[Candidate]]]
variable_input_list.clear()
value_input_list.clear()
features_list.clear()
elif 0 == tail % 4391:
process = psutil.Process()
logger.info(f"ML Validation tail {tail} meminfo:{process.memory_info()}")
if head != tail:
probability[head:tail] = self._batch_call_model(line_input_list, variable_input_list, value_input_list,
features_list)
Expand Down
9 changes: 8 additions & 1 deletion docs/source/guide.rst
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,9 @@ Get all argument list:
.. code-block:: text
usage: python -m credsweeper [-h] (--path PATH [PATH ...] | --diff_path PATH [PATH ...] | --export_config [PATH] | --export_log_config [PATH])
usage: python -m credsweeper [-h]
(--path PATH [PATH ...] | --diff_path PATH [PATH ...] | --export_config [PATH] | --export_log_config [PATH] | --git PATH [PATH ...])
[--commits POSITIVE_INT] [--branch BRANCH]
[--rules PATH] [--severity SEVERITY] [--config PATH] [--log_config PATH] [--denylist PATH]
[--find-by-ext] [--depth POSITIVE_INT] [--no-filters] [--doc] [--ml_threshold FLOAT_OR_STR]
[--ml_batch_size POSITIVE_INT] [--ml_config PATH] [--ml_model PATH] [--ml_providers STR]
Expand All @@ -31,6 +33,11 @@ Get all argument list:
exporting default config to file (default: config.json)
--export_log_config [PATH]
exporting default logger config to file (default: log.yaml)
--git PATH [PATH ...]
git repo to scan
--commits POSITIVE_INT
scan git repo for N commits only
--branch BRANCH scan git repo for single branch, otherwise - all branches were scanned (slow)
--rules PATH path of rule config file (default: credsweeper/rules/config.yaml). severity:['critical', 'high', 'medium', 'low', 'info'] type:['keyword', 'pattern', 'pem_key', 'multi']
--severity SEVERITY set minimum level for rules to apply ['critical', 'high', 'medium', 'low', 'info'](default: 'Severity.INFO', case insensitive)
--config PATH use custom config (default: built-in)
Expand Down
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ pdfminer.six==20231228
password-strength==0.0.3.post2
python-dateutil==2.8.2
pyjks==20.0.0
PyDriller==2.6
pybase62==1.0.0
base58==2.1.1

Expand Down
Loading

0 comments on commit 0894855

Please sign in to comment.