1import * as ort from "onnxruntime-web";2import type {3 FieldPrediction,4 FieldType,5 MLFieldFeatures,6} from "../shared/ml";78const LABELS: FieldType[] = ["username", "password", "email", "totp", "none"];910class BackgroundMLDetector {11 private session: ort.InferenceSession | null = null;12 private isLoading = false;13 private loadPromise: Promise<void> | null = null;1415 async init(modelUrl?: string): Promise<void> {16 if (this.session) return;17 if (this.isLoading && this.loadPromise) {18 await this.loadPromise;19 return;20 }2122 this.isLoading = true;23 this.loadPromise = this.loadModel(modelUrl);24 try {25 await this.loadPromise;26 } finally {27 this.isLoading = false;28 }29 }3031 private async loadModel(modelUrl?: string): Promise<void> {32 ort.env.wasm.wasmPaths = chrome.runtime.getURL("assets/");33 ort.env.wasm.numThreads = 1;34 ort.env.wasm.simd = true;3536 const url = modelUrl || chrome.runtime.getURL("models/form_detector.onnx");37 this.session = await ort.InferenceSession.create(url);38 }3940 async predict(features: MLFieldFeatures): Promise<FieldPrediction> {41 if (!this.session) {42 await this.init();43 }44 if (!this.session) {45 throw new Error("ML model not initialized");46 }4748 const vector = this.featuresToVector(features);49 const input = new ort.Tensor("float32", vector, [1, 45]);50 const inputName = this.session.inputNames[0];51 const outputName = this.session.outputNames[0];52 const results = await this.session.run({ [inputName]: input });53 const probs = results[outputName].data as Float32Array;5455 let maxIndex = 0;56 let maxProb = probs[0];57 for (let i = 1; i < probs.length; i += 1) {58 if (probs[i] > maxProb) {59 maxProb = probs[i];60 maxIndex = i;61 }62 }6364 return {65 type: LABELS[maxIndex],66 confidence: maxProb,67 probabilities: {68 username: probs[0],69 password: probs[1],70 email: probs[2],71 totp: probs[3],72 none: probs[4],73 },74 };75 }7677 private featuresToVector(features: MLFieldFeatures): Float32Array {78 return new Float32Array([79 features.type_text,80 features.type_email,81 features.type_password,82 features.type_tel,83 features.type_number,84 features.type_search,85 features.type_url,86 features.type_other,87 features.auto_username,88 features.auto_email,89 features.auto_current_password,90 features.auto_new_password,91 features.auto_one_time_code,92 features.auto_off,93 features.auto_other,94 features.name_has_user,95 features.name_has_login,96 features.name_has_email,97 features.name_has_pass,98 features.name_length,99 features.id_has_user,100 features.id_has_login,101 features.id_has_email,102 features.id_has_pass,103 features.id_length,104 features.placeholder_has_user,105 features.placeholder_has_email,106 features.placeholder_has_pass,107 features.placeholder_length,108 features.aria_label_has_user,109 features.aria_label_has_email,110 features.aria_label_has_pass,111 features.aria_label_length,112 features.parent_is_form,113 features.parent_is_div,114 features.parent_is_section,115 features.sibling_count,116 features.has_password_sibling,117 features.has_email_sibling,118 features.form_has_submit,119 features.form_action_has_login,120 features.is_required,121 features.has_placeholder,122 features.has_aria_label,123 features.inputmode_numeric,124 ]);125 }126}127128export const backgroundMLDetector = new BackgroundMLDetector();