import groovy.transform.Field
import com.sonymusic.ECRImage
import com.sonymusic.Utils

@Field String FINDINGS_OUTPUT_DIR = 'output'
@Field String FINDINGS_OUTPUT_FILE = "${FINDINGS_OUTPUT_DIR}/findings.txt"

def call(Map args) {
    def params = Utils.validateParams('dockerScan', args, [
        awsRegion: [type: String, required: true],
        awsRoleName: [type: String, required: false, defaultValue: 'ecr-scan-role'],
        ecrAccountId: [type: String, required: true],
        failBuild: [type: Boolean, required: false, defaultValue: true],
        imageName: [type: String, required: true],
        imageTag: [type: String, required: true],
        vulnerabilitiesToIgnore: [type: List, itemType: String, required: false, defaultValue: []],
        slackNotificationChannel: [type: String, required: false],
        scannerImage: [type: String, required: false, defaultValue: '086679231553.dkr.ecr.us-east-1.amazonaws.com/docker-image-scanner:latest'],
    ])

    dir("dockerScan/${params.imageName}") {
        def imageUrl = "${params.ecrAccountId}.dkr.ecr.${params.awsRegion}.amazonaws.com/${params.imageName}:${params.imageTag}"
        def imagePath = 'image.tar'

        def registries = [
            [accountId: params.ecrAccountId, region: params.awsRegion]
        ]
        // Add authentication for scanner image if its ECR differs from the target one
        ECRImage scannerImageProps = ECRImage.fromString(params.scannerImage)
        if (params.ecrAccountId != scannerImageProps.accountId || params.awsRegion != scannerImageProps.region) {
            registries.add([accountId: scannerImageProps.accountId, region: scannerImageProps.region])
        }

        withEcr(registries: registries) {
            sh "docker pull ${imageUrl}"
            sh "docker save ${imageUrl} -o ${imagePath} && chmod a+r ${imagePath}"
            sh "mkdir -p ${FINDINGS_OUTPUT_DIR} && chmod a+rwx ${FINDINGS_OUTPUT_DIR}"

            println("Scanning ${params.imageName}:${params.imageTag}")
            withAWS(role: params.awsRoleName, roleAccount: params.ecrAccountId, roleSessionName: params.awsRoleName, useNode: true) {
                scan(params, imagePath)
            }
        }
    }
}

def scan(Map params, String imagePath) {
    def exitCode = sh(
        label: "Scan Docker Image",
        script: """
            docker run --rm \
                --pull always \
                -v \$(pwd)/${imagePath}:/var/app/image.tar \
                -v \$(pwd)/${FINDINGS_OUTPUT_DIR}:/var/app/${FINDINGS_OUTPUT_DIR} \
                -e AWS_ACCESS_KEY_ID \
                -e AWS_SECRET_ACCESS_KEY \
                -e AWS_SESSION_TOKEN \
                -e ECR_REPOSITORY_NAME=${params.imageName} \
                -e IMAGE_TAG=${params.imageTag} \
                -e VULNERABILITIES_TO_IGNORE=${params.vulnerabilitiesToIgnore.join(',')} \
                -e FINDINGS_OUTPUT_FILE=${FINDINGS_OUTPUT_FILE} \
                ${params.scannerImage}
        """,
        returnStatus: true
    )

    handleExitCode(exitCode, params)
}

def createBuildStatusMessage(String message) {
    return "${env.JOB_NAME} - #${env.BUILD_NUMBER} ${message} (<${env.BUILD_URL}|Open>)"
}

def handleExitCode(int exitCode, Map params) {
    String errorMessage = "Docker scan failed due to vulnerabilities found for ${params.imageName}:${params.imageTag}."
    String warningMessage = "Docker scan completed successfully with warnings for non-blocking vulnerabilities for ${params.imageName}:${params.imageTag}."

    switch (exitCode) {
        case 0:
            println("Docker image scan passed successfully for ${params.imageName}:${params.imageTag}.")
            break
        case 1:
            if (params.failBuild) {
                sendSlackMessage('danger', errorMessage, params)
                error(withFindings(errorMessage))
            } else {
                sendSlackMessage('warning', warningMessage, params)
                println(warningMessage)
                markStageUnstable(warningMessage)
            }
            break
        case 2:
            sendSlackMessage('warning', warningMessage, params)
            println(warningMessage)
            markStageUnstable(warningMessage)
            break
    }
}

def sendSlackMessage(String color, String message, Map params) {
    if (params.slackNotificationChannel) {
        slackSend(
            channel: params.slackNotificationChannel,
            color: color,
            message: createBuildStatusMessage(message)
        )
    }
}

def markStageUnstable(String warningMessage) {
    catchError([
        message: warningMessage,
        buildResult: 'SUCCESS',
        stageResult: 'UNSTABLE'
    ]) {
        error(withFindings(warningMessage))
    }
}

def withFindings(String message) {
    String findings = readFile(FINDINGS_OUTPUT_FILE)?.trim()
    return findings ? "${message}\n${findings}" : message
}