package com.sonymusic

import org.junit.jupiter.api.BeforeEach
import org.junit.jupiter.api.Test

import static org.junit.jupiter.api.Assertions.assertNull

class DockerScanTest extends BaseGlobalVarTest {
    def dockerScan
    def awsRegion = 'us-east-1'
    def ecrAccountId = '123456789012'
    def awsRoleName = 'dummy'
    def imageName = 'dummy'
    def imageTag = 'latest'
    def slackNotificationChannel = '#dummy'

    @BeforeEach
    void setUp() {
        super.setUp()
        dockerScan = loadScript("vars/dockerScan.groovy")
        helper.registerAllowedMethod('sh')
        helper.registerAllowedMethod('withAWS', [Map, Closure])
        helper.registerAllowedMethod('withEcr', [Closure])
        helper.registerAllowedMethod('withEcr', [Map, Closure])
        helper.registerAllowedMethod('catchError', [Map])
        helper.registerAllowedMethod('slackSend', [Map])
        helper.registerAllowedMethod('readFile', [String.class], { String path -> '' })
    }

    @Test
    void testCallWithMissingEcrAccountId() {
        assertThrowsWithMessage(
            IllegalArgumentException.class,
            "dockerScan: Missing required parameter: 'ecrAccountId'",
            {
                dockerScan.call(
                    awsRegion: awsRegion,
                    imageName: imageName,
                    imageTag: imageTag
                )
            }
        )
    }

    @Test
    void testCallWithMissingRegion() {
        assertThrowsWithMessage(
            IllegalArgumentException.class,
            "Missing required parameter: 'awsRegion'",
            {
                dockerScan.call(
                    ecrAccountId: ecrAccountId,
                    imageName: imageName,
                    imageTag: imageTag
                )
            }
        )
    }

    @Test
    void testCallWithMissingImageName() {
        assertThrowsWithMessage(
            IllegalArgumentException.class,
            "Missing required parameter: 'imageName'",
            {
                dockerScan.call(
                    awsRegion: awsRegion,
                    ecrAccountId: ecrAccountId,
                    imageTag: imageTag
                )
            }
        )
    }

    @Test
    void testCallWithMissingImageTag() {
        assertThrowsWithMessage(
            IllegalArgumentException.class,
            "Missing required parameter: 'imageTag'",
            {
                dockerScan.call(
                    awsRegion: awsRegion,
                    ecrAccountId: ecrAccountId,
                    imageName: imageName
                )
            }
        )
    }

    @Test
    void testCallWithRequiredArgs() {
        helper.addShMock(~/(?s).*docker run.*/, 'scan successful', 0)
        dockerScan(
            awsRegion: awsRegion,
            ecrAccountId: ecrAccountId,
            imageName: imageName,
            imageTag: imageTag
        )

        assertMethodCalledOnceWith('withAWS', [
            'role=ecr-scan-role',
            "roleAccount=${ecrAccountId}",
            'roleSessionName=ecr-scan-role',
            'useNode=true'
        ])
    }

    @Test
    void testCallWithCrossAccountRole() {
        helper.addShMock(~/(?s).*docker run.*/, 'scan successful', 0)
        dockerScan(
            awsRegion: awsRegion,
            ecrAccountId: ecrAccountId,
            imageName: imageName,
            imageTag: imageTag,
            awsRoleName: awsRoleName
        )

        assertMethodCalledOnceWith('withAWS', [
            "role=${awsRoleName}",
            "roleAccount=${ecrAccountId}",
            "roleSessionName=${awsRoleName}",
            'useNode=true'
        ])
    }

    @Test
    void testCallWithCrossAccountEcr() {
        helper.addShMock(~/(?s).*docker run.*/, 'scan successful', 0)
        dockerScan(
            awsRegion: awsRegion,
            ecrAccountId: ecrAccountId,
            imageName: imageName,
            imageTag: imageTag,
            awsRoleName: awsRoleName
        )
        assertMethodCalledOnceWith('withEcr', [
            registries: [
                [accountId: ecrAccountId, region: awsRegion],
                [accountId: '086679231553', region: 'us-east-1']
            ]
        ])
    }

    @Test
    void testCallWithWrongTypeForVulnerabilitiesToIgnore() {
        assertThrowsWithMessage(
            IllegalArgumentException.class,
            "Invalid type for parameter 'vulnerabilitiesToIgnore'",
            {
                dockerScan.call(
                    awsRegion: awsRegion,
                    ecrAccountId: ecrAccountId,
                    imageName: imageName,
                    imageTag: imageTag,
                    vulnerabilitiesToIgnore: 'CVE-1'
                )
            }
        )
    }

    @Test
    void testCallWithSuccess() {
        // This one should succeed and return nothing
        assertNull(dockerScan(
            awsRegion: awsRegion,
            ecrAccountId: ecrAccountId,
            imageName: imageName,
            imageTag: imageTag,
        ))

        assertMethodCalledTimes('sh', 4)

        assertMethodCalledNthWith('sh', 0, "docker pull ${ecrAccountId}.dkr.ecr.${awsRegion}.amazonaws.com/${imageName}:latest")
        assertMethodCalledNthWith('sh', 1, "docker save ${ecrAccountId}.dkr.ecr.${awsRegion}.amazonaws.com/${imageName}:latest")
        assertMethodCalledNthWith('sh', 2, "mkdir -p output && chmod a+rwx output")
        assertMethodCalledNthWith('sh', 3, [
            'docker run',
            "-e IMAGE_TAG=${imageTag}",
            "-e ECR_REPOSITORY_NAME=${imageName}",
            "-e VULNERABILITIES_TO_IGNORE=",
            '-e FINDINGS_OUTPUT_FILE=output/findings.txt',
            'output:/var/app/output',
            'docker-image-scanner:latest'
        ])
    }

    @Test
    void testCallWithVulnerabilitiesToIgnore() {
        // This one should succeed and return nothing
        assertNull(dockerScan(
            awsRegion: awsRegion,
            ecrAccountId: ecrAccountId,
            imageName: imageName,
            imageTag: imageTag,
            vulnerabilitiesToIgnore: ['CVE-1', 'CVE-2']
        ))

        assertMethodCalledTimes('sh', 4)

        assertMethodCalledNthWith('sh', 3, [
            'docker run',
            "-e IMAGE_TAG=${imageTag}",
            "-e ECR_REPOSITORY_NAME=${imageName}",
            "-e VULNERABILITIES_TO_IGNORE=CVE-1,CVE-2",
            '-e FINDINGS_OUTPUT_FILE=output/findings.txt',
            'docker-image-scanner:latest'
        ])
    }

    @Test
    void testCallWithScanFailureFailBuildNotSet() {
        helper.registerAllowedMethod('sh', [String]) { args ->
            updateBuildStatus('FAILURE')
        }

        assertNull(dockerScan(
            awsRegion: awsRegion,
            ecrAccountId: ecrAccountId,
            imageName: imageName,
            imageTag: imageTag
        ))

        assertJobStatusFailure()

        // Verify that the catchError step was not called
        assertMethodCalledTimes('catchError', 0)
    }

    @Test
    void testCallWithScanFailureFailBuildTrue() {
        helper.addShMock(~/(?s).*docker run.*/, 'scan failed', 1)

        dockerScan(
            awsRegion: awsRegion,
            ecrAccountId: ecrAccountId,
            imageName: imageName,
            imageTag: imageTag,
            failBuild: true
        )

        assertJobStatusFailure()

        assertMethodCalledOnceWith('error', "Docker scan failed due to vulnerabilities found for ${imageName}:${imageTag}.")
    }

    @Test
    void testCallWithScanFailureIncludesFindingsInError() {
        helper.addShMock(~/(?s).*docker run.*/, 'scan failed', 1)
        helper.registerAllowedMethod('readFile', [String.class], { String path -> 'CVE-123: Critical vulnerability found' })

        dockerScan(
            awsRegion: awsRegion,
            ecrAccountId: ecrAccountId,
            imageName: imageName,
            imageTag: imageTag,
            failBuild: true
        )

        assertMethodCalledOnceWith('error', [
            "Docker scan failed due to vulnerabilities found for ${imageName}:${imageTag}.",
            "CVE-123: Critical vulnerability found"
        ])
    }

    @Test
    void testCallWithScanFailureFailBuildFalse() {
        // Mocking the shell command execution
        helper.addShMock(~/(?s).*docker run.*/, 'scan failed', 1)

        // Call dockerScan
        dockerScan(
            awsRegion: awsRegion,
            ecrAccountId: ecrAccountId,
            imageName: imageName,
            imageTag: imageTag,
            failBuild: false
        )

        assertMethodCalledOnceWith('catchError', [
            buildResult: 'SUCCESS',
            stageResult: 'UNSTABLE'
        ])
    }

    @Test
    void testCallWithSlackNotificationOnFailure() {
        // Mock the command to simulate a failure
        helper.addShMock(~/(?s).*docker run.*/, 'test failed', 1)

        dockerScan(
            awsRegion: awsRegion,
            ecrAccountId: ecrAccountId,
            imageName: imageName,
            imageTag: imageTag,
            slackNotificationChannel: slackNotificationChannel
        )

        // Verify that the slackSend method was called
        assertMethodCalledOnceWith('slackSend', [
            channel: slackNotificationChannel,
            color: 'danger',
        ])

        assertMethodCalledOnceWith('error', "Docker scan failed due to vulnerabilities found for ${imageName}:${imageTag}.")
    }

    @Test
    void testCallWithSlackNotificationOnWarnings() {
        // Mock the command to simulate warnings (exit code 2)
        helper.addShMock(~/(?s).*docker run.*/, 'warnings found', 2)

        dockerScan(
            awsRegion: awsRegion,
            ecrAccountId: ecrAccountId,
            imageName: imageName,
            imageTag: imageTag,
            slackNotificationChannel: slackNotificationChannel
        )

        // Verify that the slackSend method was called for warnings
        assertMethodCalledOnceWith('slackSend', [
            channel: slackNotificationChannel,
            color: 'warning',
        ])

        assertMethodCalledOnceWith('catchError', [
            buildResult: 'SUCCESS',
            stageResult: 'UNSTABLE'
        ])
    }

    @Test
    void testCallWithCustomScannerImage() {
        String customScannerImage = '0123456789012.dkr.ecr.us-west-2.amazonaws.com/custom-registry/scanner:v2.0'

        dockerScan(
            awsRegion: awsRegion,
            ecrAccountId: ecrAccountId,
            imageName: imageName,
            imageTag: imageTag,
            scannerImage: customScannerImage
        )

        assertMethodCalledTimes('sh', 4)

        assertMethodCalledNthWith('sh', 3, [
            'docker run',
            "-e IMAGE_TAG=${imageTag}",
            "-e ECR_REPOSITORY_NAME=${imageName}",
            "-e VULNERABILITIES_TO_IGNORE=",
            '-e FINDINGS_OUTPUT_FILE=output/findings.txt',
            customScannerImage
        ])
    }
}