AllForOne.py

#!/usr/bin/env python3
"""
AllForOne - Nuclei Template Collector
Modified by KUSH-COD3R
"""

import os
import sys
import subprocess
import shutil
import time
import requests
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, wait
import glob
from tabulate import tabulate


CONFIG = {
    'max_workers': 6,
    'timeout': 30,
    'output_folder': 'Templates',
    'temp_folder': 'TRASH',
    'repo_url': 'https://raw.githubusercontent.com/KUSH-COD3R/AllForOne/main/PleaseUpdateMe.txt'
}


def banner():
    print(r"\033[91m\033[93m  ,-.       _,---._ __  / \      _   _ _     ___              ___               ")
    print(r" /  )    .-'       `./ /   \    /_\ | | |   / __\__  _ __    /___\_ __   ___    ")
    print(r"(  (   ,'            `/    /|  //_\\| | |  / _\/ _ \| '__|  //  // '_ \ / _ \   ")
    print(r' \  `-"             \ \   / | /  _  \ | | / / | (_) | |    / \_//| | | |  __/   ')
    print(r"  `.              ,  \ \ /  | \_/ \_/_|_| \/   \___/|_|    \___/ |_| |_|\___|   ")
    print(r"   /`.          ,'-`----Y   |                                             			")
    print(r"  (            ;        |   '                                             				")
    print(r"  |  ,-.    ,-' Git-HUB |  /         Nuclei Template Collector                  ")
    print(r"  |  | (   |      BoX   | /	          - KUSH-COD3R                      ")
    print(r"  )  |  \  `.___________|/                                   				")
    print("  `--'   `--'                                               					\033[0m")


def git_clone(url, destination):
    env = os.environ.copy()
    env['GIT_TERMINAL_PROMPT'] = '0'
    try:
        result = subprocess.run(
            ['git', 'clone', '--depth', '1', url, destination],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.PIPE,
            env=env,
            timeout=CONFIG['timeout']
        )
        return result.returncode, result.stderr.decode().strip()
    except subprocess.TimeoutExpired:
        return 1, "Timeout"
    except Exception as e:
        return 1, str(e)


def generate_destination_folder(url):
    folder_name = os.path.basename(url.rstrip('.git'))
    counter = 1
    while os.path.exists(os.path.join(CONFIG['temp_folder'], folder_name)):
        folder_name = f"{os.path.basename(url.rstrip('.git'))}_{counter}"
        counter += 1
    return folder_name


def clone_repository(repo):
    if not repo.strip() or repo.startswith('#'):
        return None
    
    destination = generate_destination_folder(repo)
    return_code, error_msg = git_clone(repo, os.path.join(CONFIG['temp_folder'], destination))
    
    if return_code != 0 or 'Username' in error_msg or 'password' in error_msg.lower():
        print(f"\n\033[91mFailed to clone: {repo}\033[0m")
        return repo
    return None


def clone_repositories(file_url):
    try:
        response = requests.get(file_url, timeout=10)
        response.raise_for_status()
        repositories = [r.strip() for r in response.text.strip().split('\n') if r.strip() and not r.startswith('#')]
    except requests.exceptions.RequestException as e:
        print(f'\033[91mFailed to retrieve Repo List: {e}\033[0m')
        return

    total_repos = len(repositories)
    if total_repos == 0:
        print('\033[93mNo repositories found in list.\033[0m')
        return

    os.makedirs(CONFIG['temp_folder'], exist_ok=True)
    failed_repos = []

    print(f"\n\033[92mFound {total_repos} repositories to clone\033[0m\n")

    with ThreadPoolExecutor(max_workers=CONFIG['max_workers']) as executor:
        futures = [executor.submit(clone_repository, repo) for repo in repositories]

        with tqdm(total=total_repos, unit='repo', desc='Cloning repositories', ncols=80) as progress_bar:
            completed = 0
            while completed < total_repos:
                done, _ = wait(futures, return_when='FIRST_COMPLETED')
                completed += len(done)
                for future in done:
                    failed_repo = future.result()
                    if failed_repo:
                        failed_repos.append(failed_repo)
                    progress_bar.update(1)
                    progress = progress_bar.n / total_repos * 100
                    progress_bar.set_postfix({'Progress': f'{progress:.1f}%'})
                futures = [future for future in futures if not future.done()]

    print('\n\033[92mCloning process complete!\033[0m\n')

    if failed_repos:
        print("\033[91mFailed to clone the following repositories:\033[0m")
        for repo in failed_repos:
            print(f"  - {repo}")

    template_folder = os.path.join(os.path.dirname(os.path.abspath(__file__)), CONFIG['output_folder'])
    os.makedirs(template_folder, exist_ok=True)

    print(f"\n\033[93mMoving YAML templates to {CONFIG['output_folder']}...\033[0m")
    
    yaml_moved = 0
    for root, dirs, files in os.walk(CONFIG['temp_folder']):
        for file in files:
            if file.endswith('.yaml') or file.endswith('.yml'):
                source_path = os.path.join(root, file)
                cve_year = extract_cve_year(file)
                
                if cve_year:
                    destination_folder = os.path.join(template_folder, f"CVE-{cve_year}")
                else:
                    destination_folder = os.path.join(template_folder, "Vulnerability-Templates")
                
                os.makedirs(destination_folder, exist_ok=True)
                destination_path = os.path.join(destination_folder, file)
                
                if not os.path.exists(destination_path):
                    shutil.copy2(source_path, destination_path)
                    yaml_moved += 1

    print(f"\n\033[92mMoved {yaml_moved} YAML files\033[0m")
    print('\033[93mRemoving temporary files...\033[0m\n')
    
    try:
        shutil.rmtree(CONFIG['temp_folder'])
    except Exception as e:
        print(f"\033[91mWarning: Could not remove temp folder: {e}\033[0m")


def extract_cve_year(file_name):
    if file_name.startswith('CVE-') and len(file_name) >= 8 and file_name[4:8].isdigit():
        return file_name[4:8]
    return None


def count_yaml_files(folder):
    count = 0
    if not os.path.exists(folder):
        return 0
    for root, dirs, files in os.walk(folder):
        for file in files:
            if file.endswith('.yaml') or file.endswith('.yml'):
                count += 1
    return count


def summarize_templates():
    template_folder = os.path.join(os.path.dirname(os.path.abspath(__file__)), CONFIG['output_folder'])
    
    if not os.path.exists(template_folder):
        print("\n\033[93mNo templates found!\033[0m")
        return

    cve_folders = glob.glob(os.path.join(template_folder, 'CVE-*'))
    cve_yaml_count = sum(count_yaml_files(folder) for folder in cve_folders)

    vulnerability_templates_folder = os.path.join(template_folder, 'Vulnerability-Templates')
    vulnerability_yaml_count = count_yaml_files(vulnerability_templates_folder)

    total_yaml_count = cve_yaml_count + vulnerability_yaml_count

    data = [
        ["CVE Templates", cve_yaml_count],
        ["Other Vulnerability Templates", vulnerability_yaml_count],
        ["Total Templates", total_yaml_count]
    ]

    headers = ["Templates Type", "Templates Count"]
    table = tabulate(data, headers, tablefmt="fancy_grid")
    print(table)


def main():
    banner()
    print()
    
    clone_repositories(CONFIG['repo_url'])
    
    template_folder = os.path.join(os.path.dirname(os.path.abspath(__file__)), CONFIG['output_folder'])
    summarize_templates()

    print('\n\033[91m\033[93mPlease show your support by giving star to my GitHub repository "AllForOne".')
    print('GITHUB: https://github.com/KUSH-COD3R/AllForOne\033[0m')


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n\033[91mOperation cancelled by user.\033[0m")
        sys.exit(0)
    except Exception as e:
        print(f"\n\033[91mError: {e}\033[0m")
        sys.exit(1)