myvtm/addin/res/VTMModifyHeaders/background.js

/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. 
 *
 * @author didierfred@gmail.com
 * @version 0.4
 */


"use strict";

let config;
let started = 'on';
let debug_mode = true;
const isChrome = false;
let config_read_type = 'websocket';//local,file,websocket,http
const wsUrl = 'ws://127.0.0.1:9002?name=vtm_modify_header';
let socket = null;

loadConfigurationFromLocalStorage();

addListener();
// listen for change in configuration or start/stop
chrome.runtime.onMessage.addListener(notify);

function connectWebSocket(callback_function) {
  let socket = null; // 初始化 socket 变量
  let timeoutId = setTimeout(() => {
    socket.close();
    callback_function(null); // 超时回调 null
  }, 3000); // 5秒 = 5000 毫秒
  let connectionOpened = false; // 标记连接是否已经打开

  try {
    socket = new WebSocket(wsUrl);
  } catch (error) {
    log('WebSocket connection creation error: ' + error);
    callback_function(null); // 连接创建失败，回调 null
    return;
  }


  socket.onmessage = function (event) {
    log('Received message:' + event.data);
    clearTimeout(timeoutId); // 清除定时器
    if (socket && socket.readyState === WebSocket.OPEN) { // 检查连接状态
      socket.close();
    }
    callback_function(event.data);
    return;
  };

  socket.onopen = function () {
    log('WebSocket connection established');
    connectionOpened = true; // 设置连接已打开标志
    // 发送配置请求
    let requestData = '{"messageType":131073}';
    log('Send message:' + requestData);
    socket.send(requestData);
  };

  socket.onerror = function (error) {
    log('WebSocket connection error:' + error);
    clearTimeout(timeoutId); // 清除定时器
    if (socket && socket.readyState === WebSocket.OPEN) { // 检查连接状态
      socket.close();
    }
    callback_function(null); // 错误回调 null
  };

  socket.onclose = function (event) {
    log('WebSocket connection closed:' + event.reason);
    clearTimeout(timeoutId); // 清除定时器
    connectionOpened = false; // 重置连接已打开标志
    if (event.reason !== 'Connection timed out') { // 如果不是超时关闭
      callback_function(null); // 连接意外关闭，回调 null
    }
  };
}


function loadConfigurationFromLocalStorage() {
  if (config_read_type == 'local') {
    let headers = [];
    headers.push({ url_contains: "", action: "add", header_name: "terminalno", header_value: "7555980178", comment: "test", apply_on: "req", status: "on" });
    config = { format_version: "1.1", target_page: "", headers: headers, debug_mode: true, use_url_excepts: false };
    log("current new config" + JSON.stringify(config));


    storeInBrowserStorage({ config: JSON.stringify(config) });
    started = 'on';
  }
  else if (config_read_type == 'file') {
    log('loadConfigurationFromLocalStorage');
    const filePath = 'file:///D:/Run/runinfo/runcfg/config.json'; // 指定本地文件路径
    fetch(filePath)
      .then(response => response.json())
      .then(data => {
        console.log('Current new config:', JSON.stringify(data));
        config = data;
        storeInBrowserStorage({ config: JSON.stringify(data) }); // 如果不需要存储到浏览器存储中，可以注释掉这行
        started = 'on';
      })
      .catch(error => {
        console.error('Error reading file:', error);
      });
  }
  else if (config_read_type == 'websocket') {
    connectWebSocket(function (dataStr) {
      if (dataStr == null)
        return;
      let data = JSON.parse(dataStr);
      log("current new config:" + JSON.stringify(data));
      config = data
      storeInBrowserStorage({ config: JSON.stringify(data) });
      started = 'on';
    });
  }
}




function loadFromBrowserStorage(item, callback_function) {
  chrome.storage.local.get(item, callback_function);
}

function storeInBrowserStorage(item, callback_function) {
  chrome.storage.local.set(item, callback_function);
}


function cookie_keyvalues_set(original_cookies, key, value) {
  let new_element = " " + key + "=" + value; // not used if value is undefined. 
  let cookies_ar = original_cookies.split(";").filter(e => e.trim().length > 0);
  let selected_cookie_index = cookies_ar.findIndex(kv => kv.trim().startsWith(key + "="));
  if ((selected_cookie_index == -1) && (value != undefined)) cookies_ar.push(new_element);
  else {
    if (value === undefined)
      cookies_ar.splice(selected_cookie_index, 1);
    else
      cookies_ar.splice(selected_cookie_index, 1, new_element);
  }
  return cookies_ar.join(";");
}

function set_cookie_modify_cookie_value(original_set_cookie_header_content, key, new_value) {
  let trimmed = original_set_cookie_header_content.trimStart();
  let original_attributes = trimmed.indexOf(";") === -1 ? "" : trimmed.substring(trimmed.indexOf(";"))
  return key + "=" + new_value + original_attributes;
}


/*
* Standard function to log messages
*
*/

function log(message) {
  console.log(new Date() + " SimpleModifyHeader : " + message);
}

/*
* Rewrite the request header (add , modify or delete)
*
*/
function rewriteRequestHeader(e) {
  if (config.debug_mode) 
    log("Start modify request headers for url " + e.url + "use_url_excepts:" + config.use_url_excepts)
  for (let to_modify of config.headers) {
    if (config.debug_mode)
      log("to_modify.url_contains:" + to_modify.url_contains + " e.url:" + e.url + ", result:" + !e.url.includes(to_modify.url_contains.trim()));

    const shouldApplyModification = (to_modify, config, e) => {
      const isStatusOn = to_modify.status === "on";
      const isApplyOnReq = to_modify.apply_on === "req";
      const isUrlExcepted = config.use_url_excepts 
        ? to_modify.url_contains.length > 0 && e.url.includes(to_modify.url_contains.trim())
        : false;
      
      return isStatusOn && isApplyOnReq && !isUrlExcepted;
    };

    if (shouldApplyModification(to_modify, config, e)) {
      if (to_modify.action === "add") {
        let new_header = { "name": to_modify.header_name, "value": to_modify.header_value };
        e.requestHeaders.push(new_header);
        if (config.debug_mode) log("Add request header : name=" + to_modify.header_name +
          ",value=" + to_modify.header_value + " for url " + e.url + " url_contains " + to_modify.url_contains);
      }
      else if (to_modify.action === "modify") {
        for (let header of e.requestHeaders) {
          if (header.name.toLowerCase() === to_modify.header_name.toLowerCase()) {
            if (config.debug_mode) log("Modify request header :  name= " + to_modify.header_name +
              ",old value=" + header.value + ",new value=" + to_modify.header_value +
              " for url " + e.url);
            header.value = to_modify.header_value;
          }
        }
      }
      else if (to_modify.action === "delete") {
        let index = -1;
        for (let i = 0; i < e.requestHeaders.length; i++) {
          if (e.requestHeaders[i].name.toLowerCase() === to_modify.header_name.toLowerCase()) index = i;
        }
        if (index !== -1) {
          e.requestHeaders.splice(index, 1);
          if (config.debug_mode) log("Delete request header :  name=" + to_modify.header_name.toLowerCase() +
            " for url " + e.url);
        }
      }
      else if (to_modify.action === "cookie_add_or_modify") {
        let header_cookie = e.requestHeaders.find(header => header.name.toLowerCase() === "cookie");
        let new_cookie = cookie_keyvalues_set(header_cookie === undefined ? "" : header_cookie.value, to_modify.header_name, to_modify.header_value);
        if (header_cookie === undefined) {
          e.requestHeaders.push({ "name": "Cookie", "value": new_cookie });
          if (config.debug_mode) log("cookie_add_or_modify.req new_header : name=Cookie,value=" + new_cookie + " for url " + e.url);
        }
        else {
          header_cookie.value = new_cookie;
          if (config.debug_mode) log("cookie_add_or_modify.req modify_header : name=Cookie,value=" + new_cookie + " for url " + e.url);
        }
      }
      else if (to_modify.action === "cookie_delete") {
        let header_cookie = e.requestHeaders.find(header => header.name.toLowerCase() === "cookie");
        let new_cookie = cookie_keyvalues_set(header_cookie === undefined ? "" : header_cookie.value, to_modify.header_name, undefined);
        if (header_cookie === undefined) {
          if (config.debug_mode) log("cookie_delete.req: no cookie header found. doing nothing for url " + e.url);
        }
        else {
          header_cookie.value = new_cookie;
          if (config.debug_mode) log("cookie_delete.req modify_header : name=Cookie,value=" + new_cookie + " for url " + e.url);
        }
      }
    }
  }
  if (config.debug_mode) log("End modify request headers for url " + e.url);
  return { requestHeaders: e.requestHeaders };
}


/*
* Rewrite the response header (add , modify or delete)
*
*/
function rewriteResponseHeader(e) {
  //if (config.debug_mode) log("Start modify response headers for url " + e.url);
  for (let to_modify of config.headers) {
    if ((to_modify.status === "on") && (to_modify.apply_on === "res") && (!config.use_url_excepts || (config.use_url_excepts && e.url.includes(to_modify.url_contains.trim())))) {
      if (to_modify.action === "add") {
        let new_header = { "name": to_modify.header_name, "value": to_modify.header_value };
        e.responseHeaders.push(new_header);
        if (config.debug_mode) log("Add response header : name=" + to_modify.header_name
          + ",value=" + to_modify.header_value + " for url " + e.url);
      }
      else if (to_modify.action === "modify") {
        for (let header of e.responseHeaders) {
          if (header.name.toLowerCase() === to_modify.header_name.toLowerCase()) {
            if (config.debug_mode) log("Modify response header :  name= " + to_modify.header_name + ",old value="
              + header.value + ",new value=" + to_modify.header_value + " for url " + e.url);
            header.value = to_modify.header_value;
          }
        }
      }
      else if (to_modify.action === "delete") {
        let index = -1;
        for (let i = 0; i < e.responseHeaders.length; i++) {
          if (e.responseHeaders[i].name.toLowerCase() === to_modify.header_name.toLowerCase()) index = i;
        }
        if (index !== -1) {
          e.responseHeaders.splice(index, 1);
          if (config.debug_mode) log("Delete response header :  name=" + to_modify.header_name.toLowerCase()
            + " for url " + e.url);
        }
      }
      else if (to_modify.action === "cookie_add_or_modify") {
        let header_cookie = e.responseHeaders.find(header =>
          header.name.toLowerCase() === "set-cookie" &&
          header.value.toLowerCase().trim().startsWith(to_modify.header_name.toLowerCase() + "=")
        );
        let new_header_value = set_cookie_modify_cookie_value(header_cookie === undefined ? "" : header_cookie.value, to_modify.header_name, to_modify.header_value);
        if (header_cookie === undefined) {
          log("SimpleModifyHeaders.Warning: you're using cookie_add_or_modify in Response. While adding new cookie in response, this plugin only generates `Set-Cookie: cookie-name=cookie-value `, without ANY additional attributes. Add a `Set-Cookie` header if you need them. ");
          e.responseHeaders.push({ "name": "Set-Cookie", "value": new_header_value });
          if (config.debug_mode) log("cookie_add_or_modify.resp new_header : name=Cookie,value=" + new_header_value + " for url " + e.url);
        }
        else {
          header_cookie.value = new_header_value;
          if (config.debug_mode) log("cookie_add_or_modify.resp modify_header : name=Cookie,value=" + new_header_value + " for url " + e.url);
        }
      }
      else if (to_modify.action === "cookie_delete") {
        let index = e.responseHeaders.findIndex(header =>
          header.name.toLowerCase() === "set-cookie" &&
          header.value.toLowerCase().trim().startsWith(to_modify.header_name.toLowerCase() + "=")
        );
        if (index === -1) {
          if (config.debug_mode) log("cookie_delete.resp: no matching set-cookie header. doing nothing for url " + e.url);
        }
        else {
          e.responseHeaders.splice(index, 1);
          if (config.debug_mode) log("cookie_delete.resp delete_header : name=" + to_modify.header_name + " for url " + e.url);
        }
      }
    }
  }
  //if (config.debug_mode) log("End modify response headers for url " + e.url);
  return { responseHeaders: e.responseHeaders };
}


/*
* Listen for message form config.js
* if message is reload : reload the configuration
* if message is on : start the modify header
* if message is off : stop the modify header
*
**/
function notify(message) {
  if (message === "reload") {
    if (config.debug_mode) log("Reload configuration");
    loadFromBrowserStorage(['config'], function (result) {
      config = JSON.parse(result.config);
      if (started === "on") {
        removeListener();
        addListener();
      }
    });
  }
  else if (message === "off") {
    removeListener();
    started = "off";
    if (config.debug_mode) log("Stop modifying headers");
  }
  else if (message === "on") {
    addListener();
    started = "on";
    if (config.debug_mode) log("Start modifying headers");
  }
}

/*
* Add rewriteRequestHeader as a listener to onBeforeSendHeaders, only for the target pages.
* Add rewriteResponseHeader as a listener to onHeadersReceived, only for the target pages.
* Make it "blocking" so we can modify the headers.
*/
function addListener() {
  //return;
  let target = "<all_urls>";
  // need to had "extraHeaders" option for chrome https://developer.chrome.com/extensions/webRequest#life_cycle_footnote
  if (isChrome) {
    chrome.webRequest.onBeforeSendHeaders.addListener(rewriteRequestHeader,
      { urls: target.split(";") },
      ["blocking", "requestHeaders", "extraHeaders"]);

    chrome.webRequest.onHeadersReceived.addListener(rewriteResponseHeader,
      { urls: target.split(";") },
      ["blocking", "responseHeaders", "extraHeaders"]);
  }

  else {
    chrome.webRequest.onBeforeSendHeaders.addListener(rewriteRequestHeader,
      { urls: target.split(";") },
      ["blocking", "requestHeaders"]);
    chrome.webRequest.onHeadersReceived.addListener(rewriteResponseHeader,
      { urls: target.split(";") },
      ["blocking", "responseHeaders"]);
  }

}


/*
* Remove the two listener
*
*/
function removeListener() {
  chrome.webRequest.onBeforeSendHeaders.removeListener(rewriteRequestHeader);
  chrome.webRequest.onHeadersReceived.removeListener(rewriteResponseHeader);
}