main.py

import argparse
import sys
import git
import tempfile
import re
from collections import Counter
import os
import math
import json
import time
from tqdm import tqdm
from openai import OpenAI
from concurrent.futures import ThreadPoolExecutor, as_completed


# --------------------------------------------------------
API_KEY = os.getenv("DEEPSEEK_API_KEY")
if not API_KEY:
    sys.exit("Error: DEEPSEEK_API_KEY environment variable not set")

ENTROPY_THRESHOLD = 4.5
BINARY_EXTENSIONS = ('.png', '.jpg', '.jpeg', '.zip', '.exe', '.gif', '.pdf', '.tar', '.gz')

# --------------------------------------------------------


REGEX_PATTERNS = {
    # --- Generic High-Confidence ---
    "PRIVATE_KEY_HEADER": re.compile(r'-----BEGIN ((RSA|EC|OPENSSH|PGP|DSA) )?PRIVATE KEY-----'),
    "GENERIC_API_KEY": re.compile(r'(api_key|access_token|client_secret|bearer_token|auth_token|session_token)\s*[:=]\s*["\']?([a-zA-Z0-9_.-]{32,})["\']?', re.IGNORECASE),
    "PASSWORD_IN_URL": re.compile(r'[a-zA-Z0-9]+:\/\/[^:]+:[^@\s]+@[^\s"]+'), # Catches user:password@host
    "JWT_TOKEN": re.compile(r'ey[a-zA-Z0-9_-]{10,}\.ey[a-zA-Z0-9_-]{10,}\.[a-zA-Z0-9_-]{10,}'), # JSON Web Token

    # --- Cloud Providers (AWS) ---
    "AWS_ACCESS_KEY_ID": re.compile(r'(?<![A-Z0-9])[A-Z0-9]{20}(?![A-Z0-9])'), # AWS Access Key ID
    "AWS_SECRET_ACCESS_KEY": re.compile(r'(?<![A-Za-z0-9/+=])[A-Za-z0-9/+=]{40}(?![A-Za-z0-9/+=])'), # AWS Secret Access Key
    "AWS_MWS_AUTH_TOKEN": re.compile(r'amzn\.mws\.[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}'), # Amazon MWS Auth Token
    
    # --- Cloud Providers (Google / GCP) ---
    "GOOGLE_API_KEY": re.compile(r'AIza[0-9A-Za-z\\-_]{35}'), # Google API Key
    "GOOGLE_OAUTH_ACCESS_TOKEN": re.compile(r'ya29\.[0-9A-Za-z\-_]+'), # Google OAuth Access Token
    "GOOGLE_OAUTH_CLIENT_ID": re.compile(r'[0-9]+-[0-9a-zA-Z_]{32}\.apps\.googleusercontent\.com'), # Google OAuth Client ID
    "GOOGLE_SERVICE_ACCOUNT_KEY": re.compile(r'"type": "service_account"'), # Header of a GCP service account JSON file
    
    # --- Cloud Providers (Azure) ---
    "AZURE_STORAGE_KEY": re.compile(r'AccountKey=[a-zA-Z0-9+/=]{88}'), # Azure Storage Account Key
    "AZURE_CLIENT_SECRET": re.compile(r'client_secret["\']?\s*:\s*["\']?[a-zA-Z0-9_~\-\.]{30,}', re.IGNORECASE), # Common in Azure AD config

    # --- Git & CI/CD ---
    "GITHUB_TOKEN": re.compile(r'(gh[pousr]_[a-zA-Z0-9]{36})'), # GitHub tokens (new format)
    "GITHUB_OAUTH_TOKEN": re.compile(r'gho_[a-zA-Z0-9]{36}'), # GitHub OAuth token
    "GITLAB_PERSONAL_ACCESS_TOKEN": re.compile(r'glpat-[a-zA-Z0-9_-]{20}'), # GitLab Personal Access Token
    "HEROKU_API_KEY": re.compile(r'[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}'), # Heroku API Key (can also match UUIDs)
    "TRAVIS_CI_TOKEN": re.compile(r'travis_token\s*[:=]\s*["\']?([a-zA-Z0-9_-]{20,})["\']?', re.IGNORECASE),

    # --- Communication (Slack, Discord, Telegram) ---
    "SLACK_TOKEN": re.compile(r'(xox[pboar]?-[0-9]{12}-[0-9]{12}-[0-9]{12}-[a-z0-9]{32})'), # Slack tokens
    "SLACK_WEBHOOK_URL": re.compile(r'https://hooks\.slack\.com/services/T[a-zA-Z0-9_]{8,12}/B[a-zA-Z0-9_]{8,12}/[a-zA-Z0-9_]{24}'),
    "DISCORD_BOT_TOKEN": re.compile(r'[M-Z][a-zA-Z0-9_-]{23}\.[a-zA-Z0-9_-]{6}\.[a-zA-Z0-9_-]{27,38}'), # Discord Bot Token
    "DISCORD_WEBHOOK_URL": re.compile(r'https://discord\.com/api/webhooks/[0-9]{18,19}/[a-zA-Z0-9_-]{68}'),
    "TELEGRAM_BOT_TOKEN": re.compile(r'[0-9]{8,10}:[a-zA-Z0-9_-]{35}'), # Telegram Bot Token

    # --- Payment & SaaS ---
    "STRIPE_API_KEY": re.compile(r'(sk|pk)_(live|test)_[0-9a-zA-Z]{24,}'), # Stripe API keys
    "TWILIO_API_KEY": re.compile(r'SK[0-9a-fA-F]{32}'), # Twilio API Key
    "TWILIO_ACCOUNT_SID": re.compile(r'AC[a-f0-9]{32}'), # Twilio Account SID
    "MAILGUN_API_KEY": re.compile(r'key-[0-9a-zA-Z]{32}'), # Mailgun API Key
    "MAILCHIMP_API_KEY": re.compile(r'[0-9a-f]{32}-us[0-9]{1,2}'), # Mailchimp API Key
    "SENDGRID_API_KEY": re.compile(r'SG\.[a-zA-Z0-9_-]{22}\.[a-zA-Z0-9_-]{43}'), # SendGrid API Key
    "SHOPIFY_ACCESS_TOKEN": re.compile(r'shpat_[a-fA-F0-9]{32}'), # Shopify private app access token

    # --- Dev Tools & Artifacts ---
    "NPM_AUTH_TOKEN": re.compile(r'npm_[a-zA-Z0-9]{36}'), # NPM Auth Token
    "DOCKER_CONFIG_AUTH": re.compile(r'"auth": "[a-zA-Z0-9+/=]{50,}"'), # Docker config JSON auth
    "ARTIFACTORY_API_KEY": re.compile(r'\s*["\']?AKCp[a-zA-Z0-9]{10,}\s*["\']?'), # JFrog Artifactory API Key
    "DATADOG_API_KEY": re.compile(r'dd_api_key=["\']?([a-f0-9]{32})["\']?', re.IGNORECASE),
    "DATADOG_APP_KEY": re.compile(r'dd_app_key=["\']?([a-f0-9]{40})["\']?', re.IGNORECASE),
    "ATLASSIAN_API_TOKEN": re.compile(r'[a-zA-Z0-9]{24}'), # Atlassian API Token (Jira, Confluence) - lower confidence, can match other things
    "CLOUDINARY_URL": re.compile(r'cloudinary:\/\/[0-9]{15}:[a-zA-Z0-9_-]{27}@[a-z0-9_-]+'), # Cloudinary URL with secret

    # --- Crypto & DBs ---
    "DB_CONNECTION_STRING": re.compile(r'(postgres|mysql|mongodb|redis|amqp)s?:\/\/[^:]+:[^@\s]+@[^\s"]+', re.IGNORECASE), # Matches protocol://user:password@host
    "SSH_PRIVATE_KEY": re.compile(r'-----BEGIN (OPENSSH|DSA|RSA|EC) PRIVATE KEY-----'),
    "PGP_PRIVATE_KEY": re.compile(r'-----BEGIN PGP PRIVATE KEY BLOCK-----'),
    
    # --- Generic (Lower Confidence) ---
    "BASIC_AUTH_HEADER": re.compile(r'Authorization: Basic [a-zA-Z0-9+/=]{20,}', re.IGNORECASE), # Base64-encoded credentials
}

# --------------------------------------------------------
def parse_args():
    """Parses command-line arguments."""
    parser = argparse.ArgumentParser(
        prog="Git Vulnerability Scanner",
        description="Scans the last n commits of a git repository for possible secrets/sensitive data"
    )
    parser.add_argument('--repo', required=True, help="URL or path to a git repository")
    parser.add_argument('--n', type=int, default=10, help="Number of commits to scan (default: 10)")
    parser.add_argument('--out', default="report.json", help="Output JSON filename (default: report.json)")
    parser.add_argument('--only_sus', action='store_true', help="Only include suspicious diffs")
    parser.add_argument('--no-full-text', dest='full_text', action='store_false',
                        help="Use only suspicious lines instead of full diff text")
    parser.add_argument('--threads', type=int, default=4, help="Number of concurrent LLM calls")
    parser.set_defaults(full_text=True)
    return parser.parse_args()


# --------------------------------------------------------
def calculate_entropy(text):
    """Calculates Shannon entropy of a string."""
    if not text:
        return 0
    counts = Counter(text)
    probabilities = [count / len(text) for count in counts.values()]
    return -sum(p * math.log2(p) for p in probabilities)


# --------------------------------------------------------
def stream_repository(repo_path, max_commits):

    repo = git.Repo(repo_path)
    for commit in repo.iter_commits(max_count=max_commits):
        # Determine diffs
        if not commit.parents:
            diffs = commit.diff(git.NULL_TREE, create_patch=True)
        else:
            diffs = commit.parents[0].diff(commit, create_patch=True)

        for diff in diffs:
            # Skip binary files
            if diff.b_path and diff.b_path.lower().endswith(BINARY_EXTENSIONS):
                continue

            file_path = diff.b_path or diff.a_path
            diff_text = diff.diff.decode('utf-8', errors='ignore') if diff.diff else ""

            # Include commit summary as context
            commit_content = f"{commit.summary}\n{diff_text}"

            yield {
                "author": commit.author.name,
                "hash": commit.hexsha,
                "file_path": file_path,
                "text": commit_content
            }


# --------------------------------------------------------
PROMPT_HEADER = """You are given a git commit and you must identify lines that contain secrets or sensitive data.
You are also given a list of lines considered suspicious. You must output a JSON object with three keys:
- "file path": path of the changed file
- "commit hash": hash of the commit
- "findings": a list of sensitive items found, each formatted as:
{
  "type": "API Key",
  "line": "user_api_key = 'sk_live_123abc456def'",
  "reasoning": "The variable being initialized is called user_api_key",
  "confidence": "0.85"
}"""


# --------------------------------------------------------
def get_prompt(diff, regex_patterns, only_sus, full_text):
    """Builds an AI prompt for a single diff."""
    lines = diff['text'].split('\n')
    prompt = [PROMPT_HEADER, f"\nCommit author: {diff['author']}", f"Commit hash: {diff['hash']}", f"Commit file path: {diff['file_path']}"]
    if full_text:
        prompt.append(diff['text'])
    prompt.append("Suspicious lines:")

    suspicious = []
    for line in lines:
        if any(pattern.search(line) for pattern in regex_patterns.values()):
            suspicious.append(line)
            continue

        tokens = re.split(r'[\s"\'=:,;()\[\]{}]+', line.strip())
        if any(calculate_entropy(token) > ENTROPY_THRESHOLD for token in tokens if len(token) >= 12):
            suspicious.append(line)

    if suspicious or not only_sus:
        prompt.extend(suspicious)
        return '\n'.join(prompt)
    return None


# --------------------------------------------------------
def detect_secrets(diff_list, regex_patterns, only_sus, full_text, threads):
    """Generates prompts and sends them to DeepSeek AI in parallel."""
    client = OpenAI(api_key=API_KEY, base_url="https://api.deepseek.com")
    results = []

    def analyze_diff(diff):
        prompt = get_prompt(diff, regex_patterns, only_sus, full_text)
        if not prompt:
            return None

        for attempt in range(4):
            try:
                response = client.chat.completions.create(
                    model="deepseek-chat",
                    messages=[
                        {"role": "system", "content": "You are a helpful assistant that identifies secrets/sensitive information."},
                        {"role": "user", "content": prompt},
                    ],
                    response_format={"type": "json_object"},
                )
                content = response.choices[0].message.content
                return json.loads(content)
            except json.JSONDecodeError:
                print("API returned invalid JSON. Skipping.", file=sys.stderr)
                return None
            except Exception as e:
                if attempt == 3:
                    print(f"Prompting AI failed after retries: {e}", file=sys.stderr)
                else:
                    time.sleep(2 ** attempt)
        return None

    # --- Run in parallel ---
    with ThreadPoolExecutor(max_workers=min(threads, len(diff_list))) as executor:
        future_to_commit = {executor.submit(analyze_diff, d): d for d in diff_list}
        for future in tqdm(as_completed(future_to_commit), total=len(future_to_commit), desc="Analyzing commits"):
            result = future.result()
            if result:
                results.append(result)

    return results


# --------------------------------------------------------
def main():
    args = parse_args()
    repo_src = args.repo

    print(f"Scanning {repo_src} for last {args.n} commits...")

    try:
        # Clone remote repo if needed
        if repo_src.startswith(('http', 'git@', 'www')):
            with tempfile.TemporaryDirectory() as temp_dir:
                print(f"Cloning repository into {temp_dir} ...")
                git.Repo.clone_from(repo_src, temp_dir)
                diff_list = list(stream_repository(temp_dir, args.n))
        else:
            diff_list = list(stream_repository(repo_src, args.n))

        print(f"Retrieved {len(diff_list)} diffs, analyzing...")

        findings = detect_secrets(diff_list, REGEX_PATTERNS, args.only_sus, args.full_text, args.threads)

        with open(args.out, "w", encoding="utf-8") as f:
            json.dump(findings, f, indent=2, ensure_ascii=False)

        print(f"Report saved to {args.out} ({len(findings)} findings)")
    except git.exc.GitCommandError as e:
        sys.exit(f"Git error: {e}")
    except git.exc.InvalidGitRepositoryError:
        sys.exit(f"Error: '{repo_src}' is not a valid git repository.")
    except Exception as e:
        sys.exit(f"Unexpected error: {e}")


if __name__ == "__main__":
    main()