/*
 *  Extract images from CSS stylesheets.
 *
 *  Copyright (C) 2023-2024 Cisco Systems, Inc. and/or its affiliates. All rights reserved.
 *
 *  Authors: Micah Snyder
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License version 2 as
 *  published by the Free Software Foundation.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 *  MA 02110-1301, USA.
 */

use std::{ffi::CStr, os::raw::c_char};

use base64::{engine::general_purpose as base64_engine_standard, Engine as _};
use log::{debug, error, warn};
use unicode_segmentation::UnicodeSegmentation;

use crate::{
    scanners::magic_scan,
    sys::{cl_error_t, cl_error_t_CL_ERROR, cl_error_t_CL_SUCCESS, cli_ctx},
};

/// Error enumerates all possible errors returned by this library.
#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("Invalid format")]
    Format,

    #[error("Invalid parameter: {0}")]
    InvalidParameter(String),

    #[error("{0} parameter is NULL")]
    NullParam(&'static str),

    #[error("Failed to decode base64: {0}")]
    Base64Decode(String),
}

///
/// Extract images from HTML style CSS blocks
///
/// Process the CSS contents of an HTML <style> </style> block to find url() function
/// parameters that actually contain base64'd image contents directly and not a URL.
///
pub struct CssImageExtractor<'a> {
    remaining: &'a str,
}

impl<'a> CssImageExtractor<'a> {
    pub fn new(css: &'a str) -> Result<Self, Error> {
        Ok(Self { remaining: css })
    }

    fn next_base64_image(&mut self) -> Option<&str> {
        'outer: loop {
            // Find occurrence of "url" with
            if let Some(pos) = self.remaining.find("url") {
                (_, self.remaining) = self.remaining.split_at(pos + "url".len());
                // Found 'url'.
            } else {
                // No occurrence of "url"
                // No more 'url's.
                self.remaining = "";
                return None;
            };

            // Skip whitespace until we find '('
            for (pos, c) in self.remaining.grapheme_indices(true) {
                if c == "(" {
                    // Found left-paren.
                    (_, self.remaining) = self.remaining.split_at(pos + 1);
                    break;
                } else if c.contains(char::is_whitespace) {
                    // Skipping whitespace.
                    continue;
                } else {
                    // Missing left-paren after 'url'.
                    continue 'outer;
                }
            }

            // Find closing ')'
            let mut depth = 1;
            let mut url_parameter: Option<&str> = None;
            for (pos, c) in self.remaining.grapheme_indices(true) {
                if c == "(" {
                    // Found nested left-paren.
                    depth += 1;
                } else if c == ")" {
                    if depth > 1 {
                        // Found nested right-paren.
                        depth -= 1;
                    } else {
                        // Found right-paren.
                        let (contents, remaining) = self.remaining.split_at(pos);
                        url_parameter = Some(contents);

                        // Set the remaining bytes so we can look for more later.
                        self.remaining = remaining;
                        break;
                    }
                }
            }

            let mut url_parameter = if let Some(c) = url_parameter {
                c
            } else {
                // Never found a closing paren.  We're done here.
                self.remaining = "";
                return None;
            };

            // Strip optional whitespace and quotes from front and back.

            // Trim off whitespace at beginning
            for (pos, c) in url_parameter.grapheme_indices(true) {
                if c.contains(char::is_whitespace) {
                    // Skipping whitespace before url contents.
                    continue;
                } else {
                    (_, url_parameter) = url_parameter.split_at(pos);
                    break;
                }
            }

            // Trim off whitespace at end
            for (pos, c) in url_parameter.graphemes(true).rev().enumerate() {
                if c.contains(char::is_whitespace) {
                    // Skipping whitespace after url contents.
                    continue;
                } else {
                    (url_parameter, _) = url_parameter.split_at(url_parameter.len() - pos);
                    break;
                }
            }

            // Trim off " at beginning.
            let c = url_parameter.graphemes(true).next();
            if let Some(c) = c {
                if c == "\"" {
                    (_, url_parameter) = url_parameter.split_at(1);
                }
            };

            // Trim off " at end.
            let c = url_parameter.graphemes(true).next_back();
            if let Some(c) = c {
                if c == "\"" {
                    (url_parameter, _) = url_parameter.split_at(url_parameter.len() - 1);
                }
            };

            // Trim off whitespace at beginning.
            for (pos, c) in url_parameter.grapheme_indices(true) {
                if c.contains(char::is_whitespace) {
                    // Skipping whitespace before url contents.
                    continue;
                } else {
                    (_, url_parameter) = url_parameter.split_at(pos);
                    break;
                }
            }

            // Trim off whitespace at end.
            for (pos, c) in url_parameter.graphemes(true).rev().enumerate() {
                if c.contains(char::is_whitespace) {
                    // Skipping whitespace after url contents.
                    continue;
                } else {
                    (url_parameter, _) = url_parameter.split_at(url_parameter.len() - pos);
                    break;
                }
            }

            // Check for embedded image data for the "url"
            if !url_parameter.starts_with("data:") {
                // It's not embedded image data, let's move along.
                continue 'outer;
            }

            // Found "data:"
            (_, url_parameter) = url_parameter.split_at("data:".len());

            // The exact image type doesn't matter at all to a browser.
            // They really don't care if it's "image/gif" or "blah blah blah".
            // So we'll just skip until after the next ';'.

            // Find contents after ";"
            if let Some(pos) = url_parameter.find(';') {
                (_, url_parameter) = url_parameter.split_at(pos + ";".len());
                // Found ";"
            } else {
                // No occurrence of ";" in the url() parameter.
                continue 'outer;
            };

            // Skip whitespace until we find a 'b' (starting "base64")
            for (pos, c) in url_parameter.grapheme_indices(true) {
                if c == "b" {
                    // Found 'b'.
                    (_, url_parameter) = url_parameter.split_at(pos + 1);
                    break;
                } else if c.contains(char::is_whitespace) {
                    // Skipping whitespace.
                    continue;
                } else {
                    // Missing 'b' after ';'.
                    continue 'outer;
                }
            }

            // Verify rest of word is "ase64"
            if !url_parameter.starts_with("ase64") {
                // It's not base64 data
                continue 'outer;
            }

            // Found "base64"
            (_, url_parameter) = url_parameter.split_at("ase64".len());

            // Skip whitespace until we find ','
            for (pos, c) in url_parameter.grapheme_indices(true) {
                if c == "," {
                    // Found ','.
                    (_, url_parameter) = url_parameter.split_at(pos + 1);
                    break;
                } else if c.contains(char::is_whitespace) {
                    // Skipping whitespace.
                    continue;
                } else {
                    // Missing ',' after 'base64'.
                    continue 'outer;
                }
            }

            // Trim off whitespace at beginning.
            for (pos, c) in url_parameter.grapheme_indices(true) {
                if c.contains(char::is_whitespace) {
                    // Skipping whitespace before url contents.
                    continue;
                } else {
                    (_, url_parameter) = url_parameter.split_at(pos);
                    break;
                }
            }

            debug!("Found base64'd image data CSS url() function args.");
            return Some(url_parameter);
        }
    }
}

impl<'a> Iterator for CssImageExtractor<'a> {
    type Item = Vec<u8>;

    fn next(&mut self) -> Option<Vec<u8>> {
        // Find the next base64 encoded image
        if let Some(base64_image) = self.next_base64_image() {
            // Decode the base64 encoded image
            base64_engine_standard::STANDARD.decode(base64_image).ok()
        } else {
            None
        }
    }
}

/// C interface to handle HTML style blocks.
///
/// This function looks through the style block for images found in HTML style blocks.
/// It will extract each image and scan it.
///
/// # Safety
///
/// `file_bytes` must not be NULL
#[export_name = "html_style_block_handler"]
pub unsafe extern "C" fn html_style_block_handler(
    ctx: *mut cli_ctx,
    file_bytes: *const c_char,
) -> cl_error_t {
    let css_input = if file_bytes.is_null() {
        warn!("html css 'file_bytes' buffer is NULL");
        return cl_error_t_CL_ERROR;
    } else {
        #[allow(unused_unsafe)]
        unsafe { CStr::from_ptr(file_bytes) }.to_string_lossy()
    };

    let extractor = if let Ok(extractor) = CssImageExtractor::new(&css_input) {
        extractor
    } else {
        return cl_error_t_CL_ERROR;
    };

    let mut scan_result = cl_error_t_CL_SUCCESS;

    extractor.into_iter().all(|image| {
        debug!("Extracted {}-byte image", image.len());

        let ret = magic_scan(ctx, &image, None);
        if ret != cl_error_t_CL_SUCCESS {
            scan_result = ret;
            return false;
        }

        true
    });

    scan_result
}

#[cfg(test)]
mod tests {
    use super::*;
    use sha1::{Digest, Sha1};

    #[test]
    fn extract_data_image_gif() {
        let doc = "
            html {
                background: url(data:image/gif;base64,[encoded data, 11692 characters]);
            }
            ";

        let mut extractor = CssImageExtractor::new(&doc).unwrap();

        let image = extractor.next().unwrap();

        let mut hasher = Sha1::new();
        hasher.update(&image);
        let hash = hasher.finalize();

        assert_eq!(
            hash.as_slice(),
            [
                100, 170, 89, 45, 242, 93, 238, 12, 90, 181, 195, 223, 148, 123, 222, 106, 39, 76,
                74, 77
            ]
        )
    }

    #[test]
    fn extract_data_image_gif_in_quotes() {
        let doc = "
            html {
                background: url( \" data:image/gif;base64,[encoded data, 11692 characters] \" );
            }
            ";

        let mut extractor = CssImageExtractor::new(&doc).unwrap();

        let image = extractor.next().unwrap();

        let mut hasher = Sha1::new();
        hasher.update(&image);
        let hash = hasher.finalize();

        assert_eq!(
            hash.as_slice(),
            [
                100, 170, 89, 45, 242, 93, 238, 12, 90, 181, 195, 223, 148, 123, 222, 106, 39, 76,
                74, 77
            ]
        );
    }

    #[test]
    fn extract_data_image_gif_and_png_in_quotes() {
        let doc = "
            html {
                background: url( \" data:image/gif;base64,[encoded data, 11692 characters] \" );
                background: url( \"   data:image/png;base64,[encoded data, 14216 characters]  \"  );
            }
            ";

        let mut extractor = CssImageExtractor::new(&doc).unwrap();

        let image = extractor.next().unwrap();

        let mut hasher = Sha1::new();
        hasher.update(&image);
        let hash = hasher.finalize();

        assert_eq!(
            hash.as_slice(),
            [
                100, 170, 89, 45, 242, 93, 238, 12, 90, 181, 195, 223, 148, 123, 222, 106, 39, 76,
                74, 77
            ]
        );

        let image = extractor.next().unwrap();

        let mut hasher = Sha1::new();
        hasher.update(&image);
        let hash = hasher.finalize();

        assert_eq!(
            hash.as_slice(),
            [
                127, 44, 70, 143, 148, 237, 88, 201, 162, 82, 121, 211, 72, 66, 248, 201, 215, 6,
                242, 112
            ]
        );
    }

    #[test]
    fn extract_data_image_gif_and_png_in_quotes_2() {
        let doc = "
            url {
                background: url( \" data:image/gif;base64,[encoded data, 11692 characters] \" );
            }
            url {
                background: url( \"   data:image/png;base64,[encoded data, 14216 characters]  \"  );
            }
            ";

        let mut extractor = CssImageExtractor::new(&doc).unwrap();

        let image = extractor.next().unwrap();

        let mut hasher = Sha1::new();
        hasher.update(&image);
        let hash = hasher.finalize();

        assert_eq!(
            hash.as_slice(),
            [
                100, 170, 89, 45, 242, 93, 238, 12, 90, 181, 195, 223, 148, 123, 222, 106, 39, 76,
                74, 77
            ]
        );

        let image = extractor.next().unwrap();

        let mut hasher = Sha1::new();
        hasher.update(&image);
        let hash = hasher.finalize();

        assert_eq!(
            hash.as_slice(),
            [
                127, 44, 70, 143, 148, 237, 88, 201, 162, 82, 121, 211, 72, 66, 248, 201, 215, 6,
                242, 112
            ]
        );
    }

    #[test]
    fn extract_data_image_gif_and_png_invalid_type() {
        let doc = "
            url {
                background: url( \" data:image/gif;base64,[encoded data, 11692 characters] \" );
            }
            url {
                background: url( \"   data: blah blah blah ;  base64 , [encoded data, 14216 characters]  \"  );
            }
            ";

        let mut extractor = CssImageExtractor::new(&doc).unwrap();

        let image = extractor.next().unwrap();

        let mut hasher = Sha1::new();
        hasher.update(&image);
        let hash = hasher.finalize();

        assert_eq!(
            hash.as_slice(),
            [
                100, 170, 89, 45, 242, 93, 238, 12, 90, 181, 195, 223, 148, 123, 222, 106, 39, 76,
                74, 77
            ]
        );

        let image = extractor.next().unwrap();

        let mut hasher = Sha1::new();
        hasher.update(&image);
        let hash = hasher.finalize();

        assert_eq!(
            hash.as_slice(),
            [
                127, 44, 70, 143, 148, 237, 88, 201, 162, 82, 121, 211, 72, 66, 248, 201, 215, 6,
                242, 112
            ]
        );
    }
}
