#!/usr/bin/env python3
import struct
import base64
import json
import argparse
import sys
import os
import time

class ImaReader:
    def __init__(self, data):
        self.data = data
        self.pos = 0

    def read(self, n):
        if self.pos + n > len(self.data):
            return None
        res = self.data[self.pos:self.pos + n]
        self.pos += n
        return res

    def u32_le(self):
        b = self.read(4)
        if not b: return None
        return struct.unpack("<I", b)[0]

    def remaining(self):
        return len(self.data) - self.pos

def parse_ima_log(raw_data):
    reader = ImaReader(raw_data)
    reference_values = []

    while reader.remaining() > 0:
        pcr = reader.u32_le()
        if pcr is None: break

        pcr_digest = reader.read(32)
        if pcr_digest is None: break

        tmpl_len = reader.u32_le()
        if tmpl_len is None: break
        tmpl_raw = reader.read(tmpl_len)
        if tmpl_raw is None: break
        tmpl_name = tmpl_raw.decode('utf-8', errors='ignore').strip('\0')

        data_len = reader.u32_le()
        if data_len is None: break
        template_data = reader.read(data_len)
        if template_data is None: break

        # Filter: Support ima-ng only
        if tmpl_name == "ima-ng":
            d_reader = ImaReader(template_data)

            hash_field_len = d_reader.u32_le()
            if hash_field_len is not None:
                hash_data = d_reader.read(hash_field_len)

                if hash_data and b"sha256:" in hash_data:
                    # Extract the last 32 bytes (the SHA256 hash)
                    file_hash = hash_data[-32:].hex()

                    name_len = d_reader.u32_le()
                    if name_len is not None:
                        name_raw = d_reader.read(name_len)
                        if name_raw:
                            name = name_raw.decode('utf-8', errors='ignore').strip('\0')

                            reference_values.append({
                                "fileName": name,
                                "sha256": file_hash
                            })
    return reference_values

def main():
    parser = argparse.ArgumentParser(description="IMA Baseline Generator (ima-ng support)")
    parser.add_argument("-f", "--file", help="Input IMA log file path")
    parser.add_argument("-b", "--base64", action="store_true", help="Decode input as Base64")

    args = parser.parse_args()
    data = None

    # 1. Handling Input
    if args.file:
        if not os.path.exists(args.file):
            print(f"Error: File '{args.file}' not found.")
            sys.exit(1)
        with open(args.file, "rb") as f:
            data = f.read()
    else:
        if not sys.stdin.isatty():
            data = sys.stdin.buffer.read()
        else:
            parser.print_help()
            sys.exit(1)

    if not data:
        print("Error: No data received.")
        sys.exit(1)

    # 2. Base64 Handling
    if args.base64:
        try:
            data = base64.b64decode(data, validate=True)
        except Exception as e:
            print(f"Error: Failed to decode Base64 input - {e}")
            sys.exit(1)

    # 3. Parsing logic
    records = parse_ima_log(data)
    if not records:
        print("Error: No 'ima-ng' records found.")
        print("Hint: If the file is Base64 encoded, you MUST use the -b flag.")
        print("Hint: If the file is binary, ensure it contains 'ima-ng' templates.")
        sys.exit(1)
    result = {"referenceValues": records}

    # 4. Save to Local File
    timestamp = int(time.time())
    output_filename = f"ima_ref_value_{timestamp}.json"

    try:
        with open(output_filename, "w", encoding="utf-8") as f:
            json.dump(result, f, indent=2, ensure_ascii=False)

        # 5. Final Output
        print(f"Success: Baseline saved to '{output_filename}'")
        print(f"Total records processed: {len(records)}")
    except Exception as e:
        print(f"Error: Failed to write output file - {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()