oceanengine-dmp-pb/generate.py

# -*- coding: utf-8 -*-
from __future__ import unicode_literals
import base64
import argparse
import hashlib
from zipfile import ZipFile
import dmp_message_pb2  # 由pb文件生成的python代码, 使用Protocol Buffer2

TYPES = {
    'IMEI': dmp_message_pb2.IdItem.IMEI_MD5,
    'OAID': dmp_message_pb2.IdItem.OAID_MD5,
    'IDFA': dmp_message_pb2.IdItem.IDFA
}

def generate_valid_file(filename, type, md5):
    dataType = TYPES[type]
    target_filename = './' + type
    target_file = open(target_filename, 'w')
    count = 0
    dmp_data = dmp_message_pb2.DmpData()
    
    with open(filename) as file:
        for line in file:
            count += 1
            reset = count % 10000 == 0
            if (reset):
                binary_string = dmp_data.SerializeToString()
                result_string = base64.b64encode(binary_string).decode()
                target_file.write(result_string)
                target_file.write('\n')
                dmp_data = dmp_message_pb2.DmpData()
            if (md5):
                line = hashlib.md5(line.strip().encode('utf-8')).hexdigest()
            item = dmp_data.idList.add()
            item.dataType = dataType
            item.id = line.lower().strip()
    if count % 10000 > 0:
        binary_string = dmp_data.SerializeToString()
        result_string = base64.b64encode(binary_string).decode()
        target_file.write(result_string)
        target_file.write('\n')

    target_file.close()
    with ZipFile('./' + type + '.zip', 'w') as zip:
        zip.write(target_filename)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('filename')
    parser.add_argument('-t', '--type', required=True)
    parser.add_argument('-m', '--md5', action='store_true')
    args = parser.parse_args()
    # print(args)
    generate_valid_file(args.filename, args.type, args.md5)