2-remove_duplicated.py

import os
import shutil
import hashlib

category_map = {
	"wordpress": ["wp", "wordpress"],
	"xss": ["xss"],
	"sql_injection": ["sqli", "sql_injection"],
	"local_file_inclusion": ["lfi"],
	"remote_code_execution": ["rce"],
	"cross_site_request_forgery": ["csrf"],
	"xml_external_entity": ["xxe"],
	"cve": ["cve", "cnvd", "cnnvd", "cnnvd"],
	"open_redirect": ["redirect", "open_redirect"],
	"ssrf": ["ssrf", "server_side_request_forgery"],
	"subdomain_takeover": ["subdomain_takeover", "takeover"],
	"template_injection": ["template_injection", "ssti"],
	"crlf_injection": ["crlf_injection", "crlf"],
	"directory_listing": ["directory_listing", "traversal"],
	"exposed": ["exposed", "disclosure", "sensitive", "exposure"],
	"adobe": [
		"adobe",
		"aem",
	],
	"coldfusion": ["coldfusion", "cfm"],
	"drupal": ["drupal"],
	"joomla": ["joomla"],
	"magento": ["magento"],
	"php": ["php"],
	"airflow": ["airflow"],
	"aws": ["aws", "amazon", "ec2", "s3", "lambda", "cloudfront", "cloudfront"],
	"apache": ["apache"],
	"cpanel": ["cpanel"],
	"docker": ["docker", "container", "kubernetes"],
	"git": ["git"],
	"jenkins": ["jenkins"],
	"cisco": ["cisco"],
	"api": ["api"],
	"upload": ["upload"],
	"sensitive": ["sensitive"],
	"debug": ["debug"],
	"backup": ["backup"],
	"auth": [
		"auth",
		"login",
		"signin",
		"sign_in",
		"sign-in",
		"oauth",
		"sso",
		"register",
		"signup",
		"sign_up",
		"sign-up",
		"password",
		"pwd",
		"passwd",
		"secret",
		"token",
		"credential",
		"cred",
		"jwt",
		"cookie",
		"session",
		"remember",
		"keycloak",
		"key",
	],
	"atlassian": ["atlassian", "jira", "confluence", "bitbucket", "bamboo"],
	"config": ["config", "conf", "configuration"],
	"mysql": ["mysql", "mariadb"],
	"sql": ["sql", "database", "db"],
	"default": ["default"],
	"detect": ["detect"],
	"extract": ["extract"],
	"fuzz": ["fuzz"],
	"graphql": ["graphql"],
	"http": ["http"],
	"social": [
		"social",
		"social_media",
		"facebook",
		"twitter",
		"instagram",
		"linkedin",
	],
	"favicon": ["favicon"],
	"python": ["python", "flask", "django"],
	"ftp": ["ftp"],
	"gcloud": ["gcloud", "google_cloud", "gcp"],
	"google": ["google"],
	"graphite": ["graphite"],
	"header": ["header"],
	"injection": ["injection"],
	"ibm": ["ibm"],
	"search": ["search"],
	"ldap": ["ldap"],
	"microsoft": ["microsoft", "ms"],
	"mongodb": ["mongodb", "mongo"],
	"netlify": ["netlify"],
	"oracle": ["oracle"],
	"java": [
		"java",
		"jsp",
		"jsf",
		"j2ee",
		"j2se",
		"j2me",
		"jvm",
		"jre",
		"jdk",
		"jboss",
		"tomcat",
		"glassfish",
		"wildfly",
		"jetty",
		"websphere",
		"weblogic",
		"spring",
		"struts",
		"hibernate",
		"mybatis",
		"shiro",
	],
	"javascript": ["javascript", "js"],
	"elk": ["elk", "elasticsearch", "kibana", "logstash"],
	"kafka": ["kafka"],
	"kong": ["kong"],
	"laravel": ["laravel"],
	"nginx": ["nginx"],
	"nodejs": ["nodejs", "node", "express", "npm"],
	"perl": ["perl"],
	"postgres": ["postgres", "postgresql"],
	"rabbitmq": ["rabbitmq"],
	"redis": ["redis"],
	"ruby": ["ruby", "rails"],
	"samba": ["samba"],
	"sharepoint": ["sharepoint"],
	"smtp": ["smtp"],
	"sap": ["sap"],
	"shopify": ["shopify"],
	"ssh": ["ssh"],
	"vmware": ["vmware"],
	"web": ["web"],
}


def get_all_yaml_files(dir_path):
	all_yaml_files = {}
	for dirpath, dirs, files in os.walk(dir_path):
		dirs[:] = [d for d in dirs if d != ".git" and d != "projectdiscovery__nuclei-templates"]
		for filename in files:
			if filename.endswith(".yml") or filename.endswith(".yaml"):
				all_yaml_files[filename] = os.path.join(dirpath, filename)
	return all_yaml_files


def categorize_file(file_name, category_map):
	categories = []
	for category, keywords in category_map.items():
		if any(keyword in file_name.lower() for keyword in keywords):
			categories.append(category)
	return (
		categories if categories else ["other"]
	)


def file_hash(file_path):
	hasher = hashlib.md5()
	with open(file_path, "rb") as f:
		buf = f.read()
		hasher.update(buf)
	return hasher.hexdigest()


def copy_file_to_categories(
	file_path, base_dir, category_map, category_counts, file_hashes
):
	categories = categorize_file(os.path.basename(file_path), category_map)
	file_hash_value = file_hash(file_path)
	for category in categories:
		target_dir = os.path.join(base_dir, category)
		os.makedirs(target_dir, exist_ok=True)

		if file_hash_value not in file_hashes.get(category, set()):
			shutil.copy(
				file_path, os.path.join(target_dir, os.path.basename(file_path))
			)
			category_counts[category] = category_counts.get(category, 0) + 1
			file_hashes.setdefault(category, set()).add(file_hash_value)


def get_file_size(file_path):
	return os.path.getsize(file_path)


community_path = "clone-templates"
source_of_truth = "clone-templates/projectdiscovery__nuclei-templates"
output_path = "poc"

community = get_all_yaml_files(community_path)
nucleiTemplates = get_all_yaml_files(source_of_truth)

common_templates = set(community.keys()) & set(nucleiTemplates.keys())

category_counts = {}
file_hashes = {}

for template, community_file in community.items():
	if template in common_templates and get_file_size(community_file) == get_file_size(
		nucleiTemplates[template]
	):
		os.remove(community_file)
		continue

	copy_file_to_categories(
		community_file, output_path, category_map, category_counts, file_hashes
	)

print("各类文件数量:")
all = 0
for category, count in category_counts.items():
	all = all + count
	print(f"{category}: {count}")
print(f"all: {all}")
os.system('rm -rf clone-templates')