#!/usr/bin/env python
import re
import argparse
from termcolor import colored
from pathlib import Path
from tqdm import tqdm

"""
When I perform a web app assessment and I have access to the source code, I want to ensure that
I do not miss any API endpoints. So, this script allow me to made a diff between the list of API endpoints
I seens during the assessment against the list of API endpoints really exposed by the app based on the source code.

This script to extract all endpoints paths, from a codebase, using dedicated marker like, for example, the @Path annotation.

References:
    https://docs.oracle.com/javaee%2F7%2Fapi%2F%2F/javax/ws/rs/Path.html
    https://docs.spring.io/spring-framework/docs/current/javadoc-api/org/springframework/web/bind/annotation/package-summary.html

Dependencies:
    pip3 install tqdm termcolor
"""

EXTENSION_OF_FILE_PROCESSED = [".java"]
ENCODINGS_SEQUENCE = ["utf-8", "latin"]


def load_code(source_file):
    code = ""
    for enc in ENCODINGS_SEQUENCE:
        try:
            with source_file.open(mode="r", encoding=enc) as f:
                code = f.read()
            if len(code) > 0:
                break
        except UnicodeDecodeError:
            code = ""
            pass
    if len(code) == 0:
        raise Exception(f"Cannot load the code from the file {source_file} !")
    return code


def normalize_path(endpoints):
    endpoints_normalized = []
    for edp in endpoints:
        v = edp.replace("//", "/")
        if v not in endpoints_normalized:
            endpoints_normalized.append(v)
    return endpoints_normalized


def process_code_using_spring_annotations(source_code):
    endpoints = []
    marker = r'@(Request|Delete|Get|Patch|Post|Put)Mapping\((.*?)\)'
    regex1 = r'@(?:Request|Delete|Get|Patch|Post|Put)Mapping\(["\']?([a-z0-9_{}\/\-]*)["\']?\)'
    regex2 = r'(?:value|path)\s*=\s*["\']([a-z0-9_{}\/\-]*)["\']'
    regex3 = r'(?:value|path)\s*='
    regex4 = r'@(Delete|Get|Patch|Post|Put)Mapping[\r\n]+'
    lines_of_code = source_code.split("\n")
    base_edp_path = ""
    class_or_interface_declaration_seen = False
    for line_of_code in lines_of_code:
        if " class " in line_of_code or " interface " in line_of_code:
            class_or_interface_declaration_seen = True
        if len(re.findall(marker, line_of_code)) > 0:
            regex = regex1
            if len(re.findall(regex3, line_of_code)) > 0:
                regex = regex2
            if not class_or_interface_declaration_seen:
                base_edp_path = re.findall(regex, line_of_code, re.IGNORECASE)[0]
            else:
                edp_path = re.findall(regex, line_of_code, re.IGNORECASE)[0]
                endpoints.append(f"{base_edp_path}/{edp_path}")
    if len(endpoints) == 0 and base_edp_path != "":
        endpoints.append(base_edp_path)
    if len(re.findall(regex4, source_code)) > 0:
        endpoints.append(base_edp_path)
    return endpoints


def process_code_using_jaxrs_annotations(source_code):
    endpoints = []
    marker = "@Path("
    regex = r'@Path\(["\']([a-z0-9_{}\/\-]*)["\']\)'
    lines_of_code = source_code.split("\n")
    base_edp_path = ""
    class_or_interface_declaration_seen = False
    for line_of_code in lines_of_code:
        if " class " in line_of_code or " interface " in line_of_code:
            class_or_interface_declaration_seen = True
        if marker in line_of_code:
            if not class_or_interface_declaration_seen:
                base_edp_path = re.findall(regex, line_of_code, re.IGNORECASE)[0]
            else:
                edp_path = re.findall(regex, line_of_code, re.IGNORECASE)[0]
                endpoints.append(f"{base_edp_path}/{edp_path}")
    return endpoints


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Script to extract all endpoints paths, from a codebase, using dedicated markers.")
    parser.add_argument("-b", action="store", dest="base_folder", help="Base folder containing the source code to process", required=True)
    args = parser.parse_args()
    base_folder = args.base_folder
    print(colored(f"[+] Processing files...", "yellow"))
    global_endpoints = []
    content = list(Path(base_folder).rglob("*.*"))
    source_files = [f for f in content if f.is_file()]
    for i in tqdm(range(len(source_files))):
        source_file = source_files[i]
        if (source_file.suffix in EXTENSION_OF_FILE_PROCESSED):
            source_code = load_code(source_file)
            endpoints = process_code_using_jaxrs_annotations(source_code)
            if len(endpoints) > 0:
                global_endpoints.extend(endpoints)
            endpoints = process_code_using_spring_annotations(source_code)
            if len(endpoints) > 0:
                global_endpoints.extend(endpoints)
    print(colored(f"[+] Endpoints identified:", "yellow"))
    final_endpoints_list = normalize_path(global_endpoints)
    final_endpoints_list.sort()
    print("\n".join(final_endpoints_list))
