MitmProxy/addons/utils/Utils.py

# -*- coding: utf-8 -*-
import json
import csv
import base64
import random
import socket
from datetime import datetime, timezone, timedelta

def base64url_decode(base64_str):
    size = len(base64_str) % 4
    if size == 2:
        base64_str += '=='
    elif size == 3:
        base64_str += '='
    elif size != 0:
        raise ValueError('Invalid base64 string')
    return base64.urlsafe_b64decode(base64_str.encode('utf-8'))


def parse_jwt(jwt_token):
    jwt_token_list = jwt_token.split('.')
    header = base64url_decode(jwt_token_list[0]).decode()
    payload = base64url_decode(jwt_token_list[1]).decode()
    return {
        'header': json.loads(header),
        'payload': json.loads(payload),
        'signature': jwt_token_list[-1]
    }

def generate_random_code(n):
    range_start = 10**(n-1)
    range_end = (10**n)-1
    return str(random.randint(range_start, range_end))

def get_ip_address(domain):
    try:
        return socket.gethostbyname(domain)
    except socket.gaierror as e:
        print(f"无法解析域名 {domain}: {e}")

def read_csv(filename='results.csv'):
    data = []
    try:
        with open(filename, mode='r', encoding='utf-8') as csvfile:
            reader = csv.DictReader(csvfile);
            for row in reader:
                data.append(row);
    except IOError as e:
        print(f"read_csv error occurred: {e}");
    # print(data);
    return data;

def read_json(filename='results.json'):
    data = None;
    try:
        with open(filename, 'r', encoding='utf-8') as file:
            data = json.load(file);
    except IOError as e:
        print(f"read_json error occurred: {e}");
    # print(data);
    return data;

def seconds_to_beijing_time(seconds):
    utc_time  = datetime.fromtimestamp(seconds, tz=timezone.utc);
    beijing_time = utc_time.astimezone(timezone(timedelta(hours=8)));
    formatted_time = beijing_time.strftime('%Y-%m-%d %H:%M:%S');
    return formatted_time;

def parse_set_cookie(cookies_str):
    """
    解析HTTP响应头中的'Set-Cookie'字段，并返回一个cookie字典。
    """
    attrib_keys = ['Expires', 'Max-Age', 'Domain','Path','HttpOnly', 'SameSite'];
    parts = cookies_str.split(';');
    main_key = None;
    ret_dict = {};
    ret_data = [];
    for pstr in parts:
        pstr = pstr.strip();
        a_key = None;
        a_value = True;
        if '=' in pstr:
            a_key, a_value = pstr.split('=');
            a_key = a_key.strip();
            a_value = a_value.strip();
        else:
            a_key = pstr;
        if a_key in attrib_keys:
            if main_key in ret_dict:
                ret_dict[main_key][a_key] = a_value;
                attribs = ret_dict[main_key]['attribs'];
                attribs[a_key] = a_value;
        else:
            main_key = a_key;
            attribs = {};
            if ',' in main_key:
                keys = main_key.split(',');
                for k in keys:
                    k = k.strip();
                    if k in attrib_keys:
                        attribs[k] = True;
                    else:
                        main_key = k;
            ret_dict[main_key] = {'name': main_key, 'value':a_value, 'attribs':attribs};
            ret_data.append(ret_dict[main_key]);
    return ret_data;