kapiti 0.0.3

/// Updates DNS enums: dns-specs/*.csv and src/specs/*_generated.rs. Executed via update-specs.sh.
use std::env;
use std::fs::{self, File};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::time::Duration;
use std::vec;

use anyhow::{bail, Context, Result};
use hyper::Method;
use sha2::{Digest, Sha256};
use tracing::{self, info};

use kapiti::{fetcher::Fetcher, hyper_smol, logging};

fn main() -> Result<()> {
    logging::init_logging();

    let csv_dir = Path::new("dns-specs");
    let sections_csv_path = csv_dir.join("sections.csv");
    let sections = csv::Reader::from_path(&sections_csv_path)
        .with_context(|| format!("Failed to open {:?}", sections_csv_path))?
        .deserialize()
        .map(|v| v.expect("Failed to deserialize sections.csv entry"))
        .collect();

    match env::args().nth(1) {
        Some(val) => {
            if val == "--fetch" {
                info!("Fetching csvs (--fetch):");
                smol::block_on(download_csvs(&csv_dir, &sections))?;
            } else {
                panic!("Unrecognized arg: {}", val);
            }
        }
        None => info!("Skipping csv download (no --fetch)"),
    }

    // Output paths for generated code
    let specs_dir = Path::new("src").join("specs");

    // Generate src/specs/enums_generated.rs from CSVs in dns-specs/*
    generate_enums_rs(&specs_dir.join("enums_generated.rs"), &csv_dir, &sections)?;

    // Generate src/specs/version_generated.rs after everything else
    generate_version_rs(&specs_dir.join("version_generated.rs"), &specs_dir)?;

    Ok(())
}

#[allow(non_snake_case)]
#[derive(serde::Deserialize)]
struct SectionsRecord {
    Filename: String,
    EnumName: String,
    DataType: String,
}

/// Redownloads CSV specs from IANA. This is then consumed by this tool.
async fn download_csvs(csv_dir: &Path, sections: &Vec<SectionsRecord>) -> Result<()> {
    let fetcher = Fetcher::new(1024 * 1024, Some("text/csv".to_string()));
    let client = hyper_smol::client_system(false, Duration::from_millis(5000));

    // In theory we could fetch in parallel here, but that'd probably involve importing the full futures crate to use future::try_join_all().
    // Since this is just an occasionally run utility, we can skip that for now.
    let request_count = sections.len();
    for i in 0..request_count {
        let section = sections
            .get(i)
            .with_context(|| format!("Missing section index {}", i))?;

        let url = format!(
            "https://www.iana.org/assignments/dns-parameters/{}",
            section.Filename.as_str()
        );
        let path = csv_dir.join(&section.Filename);
        info!("- Downloading {:?} to {:?}", url, path);

        let req = fetcher
            .build_request(&Method::GET, &url)
            .expect("Failed to build HTTP request");
        let mut resp = client
            .request(req)
            .await
            .with_context(|| format!("Failed to fetch {:?}", url))?;
        let path = csv_dir.join(&section.Filename);

        info!(
            "[{}/{}] {:?} status: {}",
            i + 1,
            request_count,
            section.Filename,
            resp.status()
        );
        let mut outfile = File::create(&path.as_path())?;
        fetcher
            .write_response(&section.Filename, &mut outfile, &mut resp)
            .await?;
    }

    Ok(())
}

/// Generates enums rust code from IANA csv spec files.
fn generate_enums_rs(
    enums_rs: &Path,
    csv_dir: &Path,
    sections: &Vec<SectionsRecord>,
) -> Result<()> {
    let mut enumsfile =
        fs::File::create(&enums_rs).with_context(|| "Failed to create dns enums rust file")?;
    info!("Generating {:?}", enums_rs);
    enumsfile.write(b"// This file is autogenerated by update_specs.rs. Don't touch.\n")?;
    enumsfile.write(b"use bytecheck::CheckBytes;\n")?;
    enumsfile.write(b"use rkyv::{Archive, Deserialize, Serialize};\n")?;

    // Iterate over entries in input, open their respective files
    for section in sections {
        let enum_csv_path = csv_dir.join(&section.Filename);
        info!("- Parsing {:?} from {:?}", section.EnumName, enum_csv_path);
        let entries = match section.EnumName.as_str() {
            // Rather than trying to make everything perfectly generic,
            // just have per-file/entry tweaks/cleanup via separate functions.
            "ResourceClass" => generate_resourceclass_enum(enum_csv_path)?,
            "ResourceType" => generate_resourcetype_enum(enum_csv_path)?,
            "OpCode" => generate_opcode_enum(enum_csv_path)?,
            "ResponseCode" => generate_responsecode_enum(enum_csv_path)?,
            "OPTOptionCode" => generate_optoptioncode_enum(enum_csv_path)?,
            &_ => panic!("Unsupported enum: {}", section.EnumName),
        };

        // Write enum definition

        enumsfile.write(format!("#[repr({})]\n", section.DataType).as_bytes())?;
        enumsfile.write(
            b"#[derive(
    Archive, Deserialize, Serialize, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug,
)]\n",
        )?;
        enumsfile.write(b"#[archive_attr(derive(CheckBytes))]\n")?;
        enumsfile.write(format!("pub enum {} {{\n", section.EnumName).as_bytes())?;
        for entry in &entries {
            if entry.doc {
                enumsfile.write(format!("\n    /// {}\n", entry.comment1).as_bytes())?;
                if let Some(c2) = &entry.comment2 {
                    enumsfile.write(format!("    /// {}\n", c2).as_bytes())?;
                }
            } else {
                enumsfile.write(format!("\n    // {}\n", entry.comment1).as_bytes())?;
                if let Some(c2) = &entry.comment2 {
                    enumsfile.write(format!("    // {}\n", c2).as_bytes())?;
                }
            }
            if let Some(n) = &entry.name {
                if let Some(i) = &entry.intval {
                    enumsfile.write(format!("    {} = {},\n", n, i).as_bytes())?;
                } else {
                    panic!("Missing entry.intval for entry.name={}", n)
                }
            }
        }
        enumsfile.write(b"}\n")?;

        // Write conv helpers for enum

        // usize -> enum
        enumsfile.write(
            format!(
                "\npub fn {}_int(i: usize) -> Option<{}> {{\n",
                section.EnumName.to_lowercase(),
                section.EnumName
            )
            .as_bytes(),
        )?;
        enumsfile.write(b"    match i {\n")?;
        for entry in &entries {
            if let Some(n) = &entry.name {
                if let Some(i) = &entry.intval {
                    enumsfile.write(
                        format!("        {} => Some({}::{}),\n", i, section.EnumName, n).as_bytes(),
                    )?;
                } else {
                    panic!("Missing entry.intval for entry.name={}", n)
                }
            }
        }
        enumsfile.write(b"        _ => None,\n")?;
        enumsfile.write(b"    }\n")?;
        enumsfile.write(b"}\n")?;
    }

    Ok(())
}

struct EnumEntry {
    /// Whether to use 2 slashes (false) or 3 slashes (true) on comment1 and comment2
    doc: bool,
    /// First comment line
    comment1: String,
    /// Second comment line
    comment2: Option<String>,
    /// The name of the enum, or None (comment-only)
    name: Option<String>,
    /// The int value of the enum, ignored if name=None
    intval: Option<String>,
}

#[allow(non_snake_case)]
#[derive(serde::Deserialize)]
struct ResourceClassRecord {
    Decimal: String,
    Name: String,
    Reference: String,
}

fn generate_resourceclass_enum(enum_csv_path: PathBuf) -> Result<Vec<EnumEntry>> {
    let mut entries = Vec::new();
    for record in csv::Reader::from_path(&enum_csv_path)
        .with_context(|| format!("Failed to open enum CSV file: {:?}", enum_csv_path))?
        .deserialize()
    {
        let record: ResourceClassRecord = record
            .with_context(|| format!("Failed to read/deserialize CSV file: {:?}", enum_csv_path))?;
        let ignored = record.Decimal.contains('-') || record.Name == "Unassigned";

        // Enum comment: Name and/or reference
        let mut comment1 = record.Name.clone();
        if !record.Reference.is_empty() {
            comment1.push_str(format!(" {}", record.Reference.replace('\n', " ")).as_str());
        }

        if ignored {
            // The entry is for an unassigned value or range, list as a comment
            entries.push(EnumEntry {
                doc: !ignored,
                comment1,
                comment2: Some(record.Decimal),
                name: None,
                intval: None,
            });
        } else {
            // The entry is for a single value (which may be "reserved"...)
            let name = match record.Name.as_str() {
                // When the value is reserved, include the number in the enum name
                "Reserved" => format!("RESERVED_{}", record.Decimal),
                // Manually make these names usable
                "QCLASS NONE" => "NONE".to_string(),
                "QCLASS * (ANY)" => "ANY".to_string(),
                // All others: Use the first word, capitalized
                _ => match record.Name.split_ascii_whitespace().next() {
                    Some(first_word) => first_word.to_uppercase().replace("CLASS_", ""),
                    None => panic!("Missing name '{}' in {:?}", record.Name, enum_csv_path),
                },
            };
            entries.push(EnumEntry {
                doc: !ignored,
                comment1,
                comment2: None,
                name: Some(name),
                intval: Some(record.Decimal),
            });
        }
    }
    Ok(entries)
}

#[allow(non_snake_case)]
#[derive(serde::Deserialize)]
struct ResourceTypeRecord {
    Value: String,
    TYPE: String,
    Meaning: String,
    Reference: String,
}

fn generate_resourcetype_enum(enum_csv_path: PathBuf) -> Result<Vec<EnumEntry>> {
    let mut entries = Vec::new();
    for record in csv::Reader::from_path(&enum_csv_path)
        .with_context(|| format!("Failed to open enum CSV file: {:?}", enum_csv_path))?
        .deserialize()
    {
        let record: ResourceTypeRecord = record
            .with_context(|| format!("Failed to read/deserialize CSV file: {:?}", enum_csv_path))?;
        let ignored = record.Value.contains('-') || record.TYPE == "Unassigned";

        // Enum comment: Type or Reference+Meaning
        let mut comment1: String;
        if record.Reference.is_empty() && record.Meaning.is_empty() {
            // No additional info, just write the "TYPE" (usually something like 'Unassigned')
            comment1 = record.TYPE.clone();
        } else {
            // Instead of the TYPE, write the Meaning and/or Reference
            comment1 = String::new();
            if !record.Meaning.is_empty() {
                comment1.push_str(record.Meaning.replace('\n', " ").as_str());
            }
            if !record.Reference.is_empty() {
                if !record.Meaning.is_empty() {
                    comment1.push(' ');
                }
                comment1.push_str(record.Reference.replace('\n', " ").as_str())
            }
        }

        if ignored {
            // The entry is for an unassigned value or range, list as a comment
            entries.push(EnumEntry {
                doc: !ignored,
                comment1,
                comment2: Some(record.Value),
                name: None,
                intval: None,
            });
        } else {
            // The entry is for a single value (which may be "reserved"...)
            let name = match record.TYPE.as_str() {
                // When the value is reserved, include the number in the enum name
                "Reserved" => format!("RESERVED_{}", record.Value),
                // Manually make these names usable
                "*" => "ANY".to_string(),
                // All others: Use the type as-is
                _ => match record.Meaning.contains("OBSOLETE") {
                    true => record.TYPE.replace('-', "_") + "_OBSOLETE",
                    false => record.TYPE.replace('-', "_"),
                },
            };
            entries.push(EnumEntry {
                doc: !ignored,
                comment1,
                comment2: None,
                name: Some(name),
                intval: Some(record.Value),
            });
        }
    }
    Ok(entries)
}

#[allow(non_snake_case)]
#[derive(serde::Deserialize)]
struct OpCodeRecord {
    OpCode: String,
    Name: String,
    Reference: String,
}

fn generate_opcode_enum(enum_csv_path: PathBuf) -> Result<Vec<EnumEntry>> {
    let mut entries = Vec::new();
    for record in csv::Reader::from_path(&enum_csv_path)
        .with_context(|| format!("Failed to open enum CSV file: {:?}", enum_csv_path))?
        .deserialize()
    {
        let record: OpCodeRecord = record
            .with_context(|| format!("Failed to read/deserialize CSV file: {:?}", enum_csv_path))?;
        let ignored = record.OpCode.contains('-') || record.Name == "Unassigned";

        // Enum comment: Name and/or reference
        let mut comment1 = record.Name.clone();
        if !record.Reference.is_empty() {
            comment1.push_str(format!(" {}", record.Reference.replace('\n', " ")).as_str());
        }

        if ignored {
            // The entry is for an unassigned value or range, list as a comment
            entries.push(EnumEntry {
                doc: !ignored,
                comment1,
                comment2: Some(record.OpCode),
                name: None,
                intval: None,
            });
        } else {
            // The entry is for a single value (which may be "reserved"...)
            let name = match record.Name.as_str() {
                // Manually make these names usable
                "DNS Stateful Operations (DSO)" => "DSO".to_string(),
                "IQuery  (Inverse Query, OBSOLETE)" => "IQUERY_OBSOLETE".to_string(),
                // All others: Use the name as-is
                _ => record.Name.to_uppercase().replace("OP_", ""),
            };
            entries.push(EnumEntry {
                doc: !ignored,
                comment1,
                comment2: None,
                name: Some(name),
                intval: Some(record.OpCode),
            });
        }
    }
    Ok(entries)
}

// RCODE,Name,Description,Reference
#[allow(non_snake_case)]
#[derive(serde::Deserialize)]
struct ResponseCodeRecord {
    RCODE: String,
    Name: String,
    Description: String,
    Reference: String,
}

fn generate_responsecode_enum(enum_csv_path: PathBuf) -> Result<Vec<EnumEntry>> {
    let mut entries = Vec::new();
    for record in csv::Reader::from_path(&enum_csv_path)
        .with_context(|| format!("Failed to open enum CSV file: {:?}", enum_csv_path))?
        .deserialize()
    {
        let record: ResponseCodeRecord = record
            .with_context(|| format!("Failed to read/deserialize CSV file: {:?}", enum_csv_path))?;
        let ignored = record.RCODE.contains('-') || record.Name == "Unassigned";

        // Enum comment: Name and/or description/reference
        let mut comment1 = record.Name.clone();
        if !record.Description.is_empty() {
            comment1.push_str(format!(": {}", record.Description.replace('\n', " ")).as_str());
        }
        if !record.Reference.is_empty() {
            comment1.push_str(format!(" {}", record.Reference.replace('\n', " ")).as_str());
        }

        if ignored {
            // The entry is for an unassigned value or range, list as a comment
            entries.push(EnumEntry {
                doc: !ignored,
                comment1,
                comment2: Some(record.RCODE),
                name: None,
                intval: None,
            });
        } else if record.Description == "Not Authorized" {
            // Don't print out a duplicate RCODE=9 for both versions of "NotAuth". Let the first one handle it.
            continue;
        } else if record.Name == "BADSIG" {
            // Don't print out a duplicate RCODE=16 value for both BADVERS+BADSIG. Let BADVERS handle it.
            continue;
        } else {
            // The entry is for a single value (which may be "reserved"...)
            let name = match record.Name.as_str() {
                // Manually make these names usable
                "Reserved, can be allocated by Standards Action" => "RESERVED_65535".to_string(),
                // BADVERS and BADSIG share RCODE=16, so include both in the name. We skip BADSIG above.
                "BADVERS" => "BADVERS_BADSIG".to_string(),
                // All others: Use the name as-is
                _ => record.Name.to_uppercase(),
            };
            let comment2 = match record.Name.as_str() {
                "NotAuth" => Some(String::from("NotAuth: Not Authorized [RFC2845]")),
                "BADVERS" => Some(String::from("BADSIG: TSIG Signature Failure [RFC2845]")),
                _ => None,
            };
            entries.push(EnumEntry {
                doc: !ignored,
                comment1,
                comment2,
                name: Some(name),
                intval: Some(record.RCODE),
            });
        }
    }
    Ok(entries)
}

// RCODE,Name,Description,Reference
#[allow(non_snake_case)]
#[derive(serde::Deserialize)]
struct OPTOptionCodeRecord {
    Value: String,
    Name: String,
    Status: String,
    Reference: String,
}

fn generate_optoptioncode_enum(enum_csv_path: PathBuf) -> Result<Vec<EnumEntry>> {
    let mut entries = Vec::new();
    for record in csv::Reader::from_path(&enum_csv_path)
        .with_context(|| format!("Failed to open enum CSV file: {:?}", enum_csv_path))?
        .deserialize()
    {
        let record: OPTOptionCodeRecord = record
            .with_context(|| format!("Failed to read/deserialize CSV file: {:?}", enum_csv_path))?;
        let ignored = record.Value.contains('-') || record.Name == "Unassigned";

        // Enum comment: Name and/or status/reference
        let mut comment1 = record.Name.clone();
        if !record.Status.is_empty() {
            comment1.push_str(format!(": {}", record.Status.replace('\n', " ")).as_str());
        }
        if !record.Reference.is_empty() {
            comment1.push_str(format!(" {}", record.Reference.replace('\n', " ")).as_str());
        }

        if ignored {
            // The entry is for an unassigned value or range, list as a comment
            entries.push(EnumEntry {
                doc: !ignored,
                comment1,
                comment2: Some(record.Value),
                name: None,
                intval: None,
            });
        } else {
            // The entry is for a single value (which may be "reserved"...)
            let name = match record.Value.as_str() {
                // Override the csv-defined name of "Reserved" for the zero value
                // "EXTENDED-RCODE value 0 indicates that an unextended RCODE is in use"
                "0" => "NONE".to_string(),
                _ => match record.Name.as_str() {
                    // When the value is reserved, include the number in the enum name
                    "Reserved" => format!("RESERVED_{}", record.Value),
                    // All others: Use the type as-is (but with '-' and ' ' cleaned up)
                    _ => record
                        .Name
                        .to_uppercase()
                        .replace('-', "_")
                        .replace(' ', "_"),
                },
            };
            entries.push(EnumEntry {
                doc: !ignored,
                comment1,
                comment2: None,
                name: Some(name),
                intval: Some(record.Value),
            });
        }
    }
    Ok(entries)
}

/// Generates version hash from specs directory.
fn generate_version_rs(version_rs: &Path, specs_dir: &Path) -> Result<()> {
    let mut file_contents = vec![];
    for entry in fs::read_dir(specs_dir)? {
        let entry = entry?;
        let meta = entry.metadata()?;
        let path = entry.path();
        if meta.is_dir() {
            bail!("Unexpected directory in {:?}: {:?}", specs_dir, path);
        }
        if meta.len() > 1024 * 1024 {
            bail!("File is larger than 1MB: {:?}", path);
        }
        if path.file_name() == version_rs.file_name() {
            // Don't recursively include version.rs in hash
            continue;
        }
        let mut contents = vec![];
        let filename = path
            .file_name()
            .expect("Missing filename")
            .to_str()
            .expect("Failed to convert filename")
            .to_string();
        info!(
            "Read {} bytes from {:?}",
            File::open(path)?.read_to_end(&mut contents)?,
            filename
        );
        file_contents.push((filename, contents));
    }
    file_contents.sort_by(|(namea, _), (nameb, _)| namea.cmp(nameb));

    let mut hasher = Sha256::new();
    for (filename, contents) in file_contents {
        info!("Hashing {}", filename);
        hasher.update(contents);
    }
    let hash = format!("{:x}", hasher.finalize());
    let mut hash_trunc = hash.clone();
    hash_trunc.truncate(16);
    info!("Hash: {} => {}", hash, hash_trunc);

    let mut enumsfile =
        fs::File::create(&version_rs).with_context(|| "Failed to create version rust file")?;
    info!("Generating {:?}", version_rs);
    enumsfile.write(b"// This file is autogenerated by update_specs.rs. Don't touch.\n")?;
    enumsfile.write(format!("pub const VERSION_HASH: &str = \"{}\";\n", hash_trunc).as_bytes())?;
    Ok(())
}