GITLAB_lambda_function.py

# NOTE: RECOMMENDED APPROACH - Use AWS Secrets Manager instead of environment variables
# This is one approach and should be fine-tuned based on client requirements.

import os
import json
import base64
import requests
import boto3
import traceback
from kubernetes import client
from kubernetes.client import ApiClient, ApiException
from botocore.signers import RequestSigner
import datetime
import secrets
import re


# --------------------
# GitLab API Helper
# --------------------
def get_gitlab_api_base_url(is_enterprise):
    """
    Returns the appropriate GitLab API base URL based on the enterprise flag.
    """
    if is_enterprise:
        # Get enterprise URL from environment variables
        base_url = os.environ.get("GITLAB_ENTERPRISE_URL")
        if not base_url:
            raise ValueError("is_enterprise is true, but GITLAB_ENTERPRISE_URL environment variable is not set.")
        # Expecting a URL like https://your-gitlab-instance.com/api/v4
        return base_url.rstrip('/')
    else:
        return "https://gitlab.com/api/v4"

# --------------------
# GitLab Repo Listing
# --------------------
def list_repositories(name, target_type, token, page=1, is_enterprise=False):
    """
    Fetches a list of repositories for a given GitLab user or group,
    supporting both public and enterprise GitLab.
    For users, automatically fetches both owned and contributed projects.
    """
    print(f"[INFO] Fetching GitLab repositories (Enterprise: {is_enterprise})...")
    headers = {
        "PRIVATE-TOKEN": token,
        "Accept": "application/json"
    }
    
    base_url = get_gitlab_api_base_url(is_enterprise)
    
    if target_type == "user":
        # For user repos, fetch both owned and contributed projects
        
        # 1. Get owned projects - need to get user ID first
        user_url = f"{base_url}/users?username={name}"
        user_response = requests.get(user_url, headers=headers)
        if user_response.status_code != 200 or not user_response.json():
            print(f"[ERROR] GitLab API error fetching user: {user_response.text}")
            raise Exception(f"GitLab API error fetching user: {user_response.text}")
        
        user_id = user_response.json()[0]['id']
        
        # Fetch owned projects
        owned_repos_url = f"{base_url}/users/{user_id}/projects?per_page=100&page={page}"
        print(f"[DEBUG] Fetching owned projects: {owned_repos_url}")
        owned_response = requests.get(owned_repos_url, headers=headers)
        
        if owned_response.status_code != 200:
            print(f"[ERROR] GitLab API error fetching owned repos ({owned_response.status_code}): {owned_response.text}")
            raise Exception(f"GitLab API error ({owned_response.status_code}): {owned_response.text}")
        
        owned_has_next = owned_response.headers.get("X-Next-Page", "").strip() != ""
        owned_repos = owned_response.json()
        
        # 2. Fetch contributed projects (where user is member but not owner)
        contributed_repos_url = f"{base_url}/projects?membership=true&owned=false&per_page=100&page={page}"
        print(f"[DEBUG] Fetching contributed projects: {contributed_repos_url}")
        contributed_response = requests.get(contributed_repos_url, headers=headers)
        
        contributed_repos = []
        contributed_has_next = False
        
        if contributed_response.status_code == 200:
            contributed_repos = contributed_response.json()
            contributed_has_next = contributed_response.headers.get("X-Next-Page", "").strip() != ""
            print(f"[INFO] Retrieved {len(contributed_repos)} contributed repositories")
        else:
            # Log warning but don't fail - contributed projects are supplementary
            print(f"[WARN] Could not fetch contributed projects ({contributed_response.status_code}): {contributed_response.text}")
        
        # Combine and deduplicate by project path
        all_repos = owned_repos + contributed_repos
        seen_paths = set()
        unique_repos = []
        for repo in all_repos:
            path = repo.get("path_with_namespace")
            if path not in seen_paths:
                seen_paths.add(path)
                unique_repos.append(repo)
        
        # hasNext is true if either source has more pages
        has_next = owned_has_next or contributed_has_next
        
        formatted_repos = []
        for repo in unique_repos:
            formatted_repos.append({
                "name": repo.get("path_with_namespace"),
                "cloneUrl": repo.get("http_url_to_repo"),
                "htmlUrl": repo.get("web_url"),
                "language": "Gitlab",
                "private": repo.get("visibility") == "private"
            })
        
        print(f"[INFO] Retrieved {len(formatted_repos)} total repositories (owned + contributed, deduplicated)")
        return {"content": formatted_repos, "hasNext": has_next}
    
    else:
        # For groups (orgs in GitLab)
        repos_url = f"{base_url}/groups/{name}/projects?per_page=100&page={page}&include_subgroups=true"
        
        # Debug line to confirm the final URL
        print(f"[DEBUG] Calling URL: {repos_url}")
        
        response = requests.get(repos_url, headers=headers)

        if response.status_code != 200:
            print(f"[ERROR] GitLab API error ({response.status_code}): {response.text}")
            raise Exception(f"GitLab API error ({response.status_code}): {response.text}")

        # Check for pagination
        has_next = response.headers.get("X-Next-Page", "").strip() != ""
        
        formatted_repos = []
        for repo in response.json():
            formatted_repos.append({
                "name": repo.get("path_with_namespace"),
                "cloneUrl": repo.get("http_url_to_repo"),
                "htmlUrl": repo.get("web_url"),
                "language": "Gitlab",
                "private": repo.get("visibility") == "private"
            })

        print(f"[INFO] Retrieved {len(formatted_repos)} repositories")
        return {"content": formatted_repos, "hasNext": has_next}

# --------------------
# GitLab Branch Listing
# --------------------
def list_branches(repo_url, token, is_enterprise=False):
    """
    Fetches all branch names for a given GitLab repository URL,
    supporting both public and enterprise GitLab.
    """
    print(f"[INFO] Fetching branches for repo: {repo_url} (Enterprise: {is_enterprise})")

    # Extract project path from GitLab URL
    # GitLab URLs can be like: https://gitlab.com/namespace/project.git or https://gitlab.com/namespace/subnamespace/project.git
    match = re.search(r"https://[^/]+/(.+?)(?:\.git)?$", repo_url)
    if not match:
        raise ValueError("Invalid GitLab repository URL format.")

    project_path = match.group(1)
    # URL encode the project path for API calls
    encoded_project_path = requests.utils.quote(project_path, safe='')

    headers = {
        "PRIVATE-TOKEN": token,
        "Accept": "application/json"
    }
    
    base_url = get_gitlab_api_base_url(is_enterprise)
    branches = []
    page = 1
    
    while True:
        branches_url = f"{base_url}/projects/{encoded_project_path}/repository/branches?per_page=100&page={page}"
        response = requests.get(branches_url, headers=headers)

        if response.status_code != 200:
            print(f"[ERROR] GitLab API error ({response.status_code}): {response.text}")
            raise Exception(f"GitLab API error ({response.status_code}): {response.text}")

        data = response.json()
        if not data:
            break

        for branch in data:
            branches.append(branch.get("name"))
        
        # Check if there's a next page
        if not response.headers.get("X-Next-Page", "").strip():
            break
        
        page += 1

    print(f"[INFO] Retrieved {len(branches)} branches.")
    return {"content": branches}

# --------------------
# GitLab Default Branch
# --------------------
def get_default_branch(repo_url, token, is_enterprise=False):
    """
    Fetches the default branch for a given GitLab repository URL,
    supporting both public and enterprise GitLab.
    """
    print(f"[INFO] Fetching default branch for repo: {repo_url} (Enterprise: {is_enterprise})")

    # Extract project path from GitLab URL
    match = re.search(r"https://[^/]+/(.+?)(?:\.git)?$", repo_url)
    if not match:
        raise ValueError("Invalid GitLab repository URL format.")

    project_path = match.group(1)
    # URL encode the project path for API calls
    encoded_project_path = requests.utils.quote(project_path, safe='')

    headers = {
        "PRIVATE-TOKEN": token,
        "Accept": "application/json"
    }
    
    base_url = get_gitlab_api_base_url(is_enterprise)
    repo_api_url = f"{base_url}/projects/{encoded_project_path}"
    response = requests.get(repo_api_url, headers=headers)

    if response.status_code != 200:
        print(f"[ERROR] GitLab API error ({response.status_code}): {response.text}")
        raise Exception(f"GitLab API error ({response.status_code}): {response.text}")

    default_branch = response.json().get("default_branch")
    print(f"[INFO] Retrieved default branch: {default_branch}")
    return {"content": default_branch}


# --------------------
# EKS Token Generator
# --------------------
def get_eks_token(cluster_name, region):
    """
    Generates a presigned STS URL for EKS authentication.
    """
    print("[INFO] Generating EKS token with custom signing...")
    try:
        session = boto3.session.Session()
        sts_client = session.client('sts', region_name=region)
        service_model = sts_client.meta.service_model

        signer = RequestSigner(
            service_model.service_id, region, 'sts', 'v4',
            session.get_credentials(), session.events
        )
        request_dict = {
            'method': 'GET',
            'url': f'https://sts.{region}.amazonaws.com/?Action=GetCallerIdentity&Version=2011-06-15',
            'body': {}, 'headers': {'x-k8s-aws-id': cluster_name}, 'context': {}
        }
        presigned_url = signer.generate_presigned_url(
            request_dict=request_dict, expires_in=60, operation_name='GetCallerIdentity'
        )
        token = 'k8s-aws-v1.' + base64.urlsafe_b64encode(
            presigned_url.encode('utf-8')
        ).decode('utf-8').rstrip('=')

        print("[INFO] EKS token generated successfully.")
        return token
    except Exception as e:
        print(f"[ERROR] Failed to generate EKS token for cluster '{cluster_name}'.")
        raise e

# --------------------
# EKS Job Launcher
# --------------------
def trigger_eks_job(cluster_name, region, job_name, image, env_vars, args_list, namespace="default"):
    """
    Configures and launches a job on the specified EKS cluster.
    """
    print(f"[INFO] Starting job creation for: {job_name}")
    eks = boto3.client('eks', region_name=region)
    cluster_info = eks.describe_cluster(name=cluster_name)['cluster']

    configuration = client.Configuration()
    configuration.host = cluster_info['endpoint']

    ca_data = cluster_info['certificateAuthority']['data']
    ca_path = '/tmp/ca.crt'
    with open(ca_path, 'w') as f:
        f.write(base64.b64decode(ca_data).decode())
    configuration.ssl_ca_cert = ca_path

    token = get_eks_token(cluster_name, region)
    configuration.api_key = {"authorization": f"Bearer {token}"}

    api_client = ApiClient(configuration=configuration)
    batch_v1 = client.BatchV1Api(api_client)

    print("[INFO] Defining Kubernetes job spec...")
    container = client.V1Container(
        name="scanner", image=image, args=args_list,
        env=[client.V1EnvVar(name=k, value=v) for k, v in env_vars.items()]
    )
    
    # Get the image pull secret name from environment variable
    image_pull_secret_name = os.environ.get("IMAGE_PULL_SECRET_NAME", "image-secret")
    
    pod_spec = client.V1PodSpec(
        restart_policy="Never", 
        containers=[container],
        image_pull_secrets=[client.V1LocalObjectReference(name=image_pull_secret_name)]
    )
    template = client.V1PodTemplateSpec(
        metadata=client.V1ObjectMeta(name=job_name), spec=pod_spec
    )
    spec = client.V1JobSpec(
        template=template, backoff_limit=1, ttl_seconds_after_finished=180
    )
    job = client.V1Job(
        api_version="batch/v1", kind="Job",
        metadata=client.V1ObjectMeta(name=job_name), spec=spec
    )

    print(f"[INFO] Submitting job to Kubernetes cluster with image pull secret: {image_pull_secret_name}")
    batch_v1.create_namespaced_job(namespace=namespace, body=job)
    print(f"[SUCCESS] Job '{job_name}' successfully created in namespace '{namespace}'.")

# --------------------
# EKS Job Status Checker
# --------------------
def check_job_status(cluster_name, region, job_name, namespace="default"):
    """
    Checks the status and fetches logs for a given Kubernetes job.
    """
    print(f"[INFO] Checking status for job: {job_name}")
    eks = boto3.client('eks', region_name=region)
    try:
        cluster_info = eks.describe_cluster(name=cluster_name)['cluster']
    except eks.exceptions.ResourceNotFoundException:
        print(f"[ERROR] EKS cluster '{cluster_name}' not found.")
        raise ValueError(f"EKS cluster '{cluster_name}' not found.")

    configuration = client.Configuration()
    configuration.host = cluster_info['endpoint']
    ca_data = cluster_info['certificateAuthority']['data']
    ca_path = f'/tmp/{secrets.token_hex(8)}_ca.crt'
    with open(ca_path, 'w') as f:
        f.write(base64.b64decode(ca_data).decode())
    configuration.ssl_ca_cert = ca_path
    
    token = get_eks_token(cluster_name, region)
    configuration.api_key = {"authorization": f"Bearer {token}"}
    
    api_client = ApiClient(configuration=configuration)
    batch_v1 = client.BatchV1Api(api_client)
    core_v1 = client.CoreV1Api(api_client)

    try:
        job = batch_v1.read_namespaced_job(name=job_name, namespace=namespace)
        status = "SCHEDULED"
        logs = ""
        error_message = ""

        if job.status.succeeded:
            status = "SUCCEEDED"
        elif job.status.failed:
            status = "FAILED"
        
        if status != "SCHEDULED" or (job.status.active and job.status.active > 0):
            pod_label_selector = f"job-name={job_name}"
            pod_list = core_v1.list_namespaced_pod(namespace=namespace, label_selector=pod_label_selector)
            
            if pod_list.items:
                pod_name = pod_list.items[0].metadata.name
                try:
                    logs = core_v1.read_namespaced_pod_log(name=pod_name, namespace=namespace, container="scanner")
                    if status == "FAILED":
                        error_message = "Job failed. See logs for details."
                    elif status == "SUCCEEDED":
                        logs = f"successful job completed, logs: \n{logs}"
                except ApiException as e:
                    log_error_msg = f"Could not retrieve logs for pod {pod_name}: {e.reason}"
                    print(f"[ERROR] {log_error_msg}")
                    logs = log_error_msg
                    if status == "FAILED":
                        error_message = log_error_msg
            else:
                logs = "Pod not found or not ready. Logs are not yet available."
                if status == "FAILED":
                    error_message = "Job failed, but its pod could not be found to retrieve logs."
        
        return {"status": status, "logs": logs, "error": error_message}

    except ApiException as e:
        if e.status == 404:
            print(f"[ERROR] Job '{job_name}' not found in namespace '{namespace}'.")
            return {"status": "NOT_FOUND", "logs": "", "error": f"Job '{job_name}' not found."}
        else:
            print(f"[ERROR] Kubernetes API error: {e}")
            raise e

# --------------------
# Lambda Handler
# --------------------
def lambda_handler(event, context):
    """
    Main handler for the Lambda function. Routes actions for listing repositories,
    branches, or triggering scans.
    """
    print("[INFO] Lambda triggered with event:")
    

    if "body" in event and isinstance(event["body"], str):
        try:
            event = json.loads(event["body"])
            print("[INFO] Parsed request body from Function URL.")
        except Exception as parse_err:
            print(f"[ERROR] Failed to parse 'body': {parse_err}")
            return {"statusCode": 400, "body": json.dumps("Invalid request body")}

    action = event.get("action")
    gitlab_token = os.environ.get("GITLAB_TOKEN")

    if not gitlab_token:
        error_msg = "FATAL: GITLAB_TOKEN environment variable not set."
        print(f"[ERROR] {error_msg}")
        return {"statusCode": 500, "body": json.dumps(error_msg)}

    try:
        if action == "list_repos":
            print("[INFO] Action: list_repos")
            name = event.get("name") or os.environ.get("GITLAB_USERNAME")
            target_type = event.get("type", "user")
            page = int(event.get("page", 1))
            is_enterprise = str(event.get("is-enterprise", "false")).lower() == "true"

            if not name:
                return {"statusCode": 400, "body": json.dumps("Missing required parameter: 'name'")}
            if target_type not in ["user", "org"]:
                return {"statusCode": 400, "body": json.dumps("Invalid 'type'. Must be 'user' or 'org'.")}

            repos_data = list_repositories(name, target_type, gitlab_token, page, is_enterprise=is_enterprise)
            return {"statusCode": 200, "body": json.dumps(repos_data)}

        elif action == "list_branches":
            print("[INFO] Action: list_branches")
            repo_url = event.get("repo_url")
            repo_type = event.get("type")
            is_enterprise = str(event.get("is-enterprise", "false")).lower() == "true"

            if not repo_url:
                return {"statusCode": 400, "body": json.dumps("Missing required parameter: 'repo_url'")}
            if repo_type != "GITLAB":
                return {"statusCode": 400, "body": json.dumps("Invalid 'type'. Must be 'GITLAB'.")}

            branch_data = list_branches(repo_url, gitlab_token, is_enterprise=is_enterprise)
            return {"statusCode": 200, "body": json.dumps(branch_data)}

        elif action == "default_branch":
            print("[INFO] Action: default_branch")
            repo_url = event.get("repo_url")
            repo_type = event.get("type")
            is_enterprise = str(event.get("is-enterprise", "false")).lower() == "true"

            if not repo_url:
                return {"statusCode": 400, "body": json.dumps("Missing required parameter: 'repo_url'")}
            if repo_type != "GITLAB":
                return {"statusCode": 400, "body": json.dumps("Invalid 'type'. Must be 'GITLAB'.")}

            default_branch_data = get_default_branch(repo_url, gitlab_token, is_enterprise=is_enterprise)
            return {"statusCode": 200, "body": json.dumps(default_branch_data)}

        elif action == "trigger_scan":
            print("[INFO] Action: trigger_scan")
            
            args_dict = event.get("args", {})
            if not args_dict:
                return {"statusCode": 400, "body": json.dumps("Job arguments ('args') dictionary cannot be empty.")}
            
            # Get job-id from the nested 'args' dictionary.
            job_id = args_dict.get("job-id")
            if not job_id:
                return {"statusCode": 400, "body": json.dumps("Missing required parameter 'job-id' in 'args'.")}

            # Construct the job name using the client-provided job-id.
            job_name = f"cdefense-hbyrid-{job_id}"
            print(f"[INFO] Constructed job name from client-provided id: {job_name}")

            image = os.environ.get("CLI_IMAGE")
            cluster_name = os.environ.get("EKS_CLUSTER_NAME")
            region = os.environ.get("AWS_REGION")

            if not cluster_name:
                return {"statusCode": 500, "body": json.dumps("EKS_CLUSTER_NAME environment variable not set.")}

            env_vars = event.get("env", {})
            
            if args_dict.get("repo-type") in ["GITLAB", "BRANCH"]:
                repo_url = args_dict.get("repo-url")
                if repo_url:
                    # For GitLab, we need to handle authentication differently
                    # GitLab supports token in URL like: https://oauth2:TOKEN@gitlab.com/user/repo.git
                    url_without_scheme = repo_url.replace("https://", "").replace("http://", "")
                    authenticated_url = f"https://oauth2:{gitlab_token}@{url_without_scheme}"
                    env_vars["GIT_REPO"] = authenticated_url
                    print("[INFO] Constructed authenticated GIT_REPO URL for private repository.")

            env_vars["GITLAB_TOKEN"] = gitlab_token
            env_vars["AWS_ACCESS_KEY_ID"] = os.environ.get("AWS_ACCESS_KEY_ID_HYBRID")
            env_vars["AWS_SECRET_ACCESS_KEY"] = os.environ.get("AWS_SECRET_ACCESS_KEY_HYBRID")
            env_vars["BUCKET_NAME"] = os.environ.get("BUCKET_NAME_HYBRID")
            env_vars["AWS_REGION"] = os.environ.get("AWS_REGION_HYBRID")

            args_list = ["full"]
            
            if str(args_dict.pop("is-enterprise", "false")).lower() == "true":
                args_list.append("--is-enterprise")
            
            for key, value in args_dict.items():
                args_list.append(f"--{key}={str(value)}")

            print(f"[INFO] Constructed job arguments: {args_list}")
            print("[INFO] Triggering EKS job...")
            trigger_eks_job(cluster_name, region, job_name, image, env_vars, args_list)
            print("[SUCCESS] Job submission process completed.")
            return {"statusCode": 202, "body": json.dumps(f"K8s job '{job_name}' accepted for processing.")}

        elif action == "trigger_scan_public":
            print("[INFO] Action: trigger_scan_public")
            
            args_dict = event.get("args", {})
            if not args_dict:
                return {"statusCode": 400, "body": json.dumps("Job arguments ('args') dictionary cannot be empty.")}
            
            # Get job-id from the nested 'args' dictionary.
            job_id = args_dict.get("job-id")
            if not job_id:
                return {"statusCode": 400, "body": json.dumps("Missing required parameter 'job-id' in 'args'.")}

            # Validate repo-url is provided for public repo scan
            repo_url = args_dict.get("repo-url")
            if not repo_url:
                return {"statusCode": 400, "body": json.dumps("Missing required parameter 'repo-url' for public repository scan.")}

            # Construct the job name using the client-provided job-id.
            job_name = f"cdefense-public-{job_id}"
            print(f"[INFO] Constructed job name for public repo scan: {job_name}")

            image = os.environ.get("CLI_IMAGE")
            cluster_name = os.environ.get("EKS_CLUSTER_NAME")
            region = os.environ.get("AWS_REGION")

            if not cluster_name:
                return {"statusCode": 500, "body": json.dumps("EKS_CLUSTER_NAME environment variable not set.")}

            # Set up environment variables for public repo scan
            env_vars = event.get("env", {})
            
            # For public repositories, we can use the URL directly without authentication
            env_vars["GIT_REPO"] = repo_url
            print(f"[INFO] Using public repository URL: {repo_url}")

            # Add AWS credentials for results storage
            env_vars["AWS_ACCESS_KEY_ID"] = os.environ.get("AWS_ACCESS_KEY_ID_HYBRID")
            env_vars["AWS_SECRET_ACCESS_KEY"] = os.environ.get("AWS_SECRET_ACCESS_KEY_HYBRID")
            env_vars["BUCKET_NAME"] = os.environ.get("BUCKET_NAME_HYBRID")
            env_vars["AWS_REGION"] = os.environ.get("AWS_REGION_HYBRID")

            args_list = ["full"]
            
            # Handle enterprise flag if present
            if str(args_dict.pop("is-enterprise", "false")).lower() == "true":
                args_list.append("--is-enterprise")
            
            # Add all other arguments
            for key, value in args_dict.items():
                args_list.append(f"--{key}={str(value)}")

            print(f"[INFO] Constructed job arguments for public repo: {args_list}")
            print("[INFO] Triggering EKS job for public repository scan...")
            trigger_eks_job(cluster_name, region, job_name, image, env_vars, args_list)
            print("[SUCCESS] Public repository scan job submission completed.")
            return {"statusCode": 202, "body": json.dumps(f"K8s job '{job_name}' accepted for processing.")}

        elif action == "job_status_check":
            print("[INFO] Action: job_status_check")
            job_id = event.get("job-id")
            job_type = event.get("job-type", "hybrid")  # Default to hybrid for backward compatibility
            
            if not job_id:
                return {"statusCode": 400, "body": json.dumps("Missing required parameter: 'job-id'")}
            
            # Construct the full job name to check based on job type
            if job_type == "public":
                job_name_to_check = f"cdefense-public-{job_id}"
            else:
                job_name_to_check = f"cdefense-hbyrid-{job_id}"  # Keep the original spelling for backward compatibility
            
            print(f"[INFO] Checking status for job name: {job_name_to_check} (type: {job_type})")

            cluster_name = os.environ.get("EKS_CLUSTER_NAME")
            region = os.environ.get("AWS_REGION")

            if not cluster_name:
                return {"statusCode": 500, "body": json.dumps("EKS_CLUSTER_NAME environment variable not set.")}

            status_data = check_job_status(cluster_name, region, job_name_to_check)
            return {"statusCode": 200, "body": json.dumps(status_data)}

        else:
            print(f"[ERROR] Invalid action received: {action}")
            return {"statusCode": 400, "body": json.dumps("Invalid or missing 'action'.")}

    except Exception as e:
        error_trace = traceback.format_exc()
        print(f"[EXCEPTION] An unexpected error occurred: {str(e)}")
        print("[TRACEBACK]", error_trace)
        return {
            "statusCode": 500,
            "body": json.dumps({
                "error": "An internal server error occurred.",
                "details": str(e),
                "traceback": error_trace
            })
        }