# repack_netflow_ipk.py — 将修改后的 control/ 和 data/ 目录重新打包为 IPK (tar.gz 格式)
# OpenWrt IPK 格式: 外层 tar.gz，内含 debian-binary + control.tar.gz + data.tar.gz
# 用法: python repack_netflow_ipk.py <data_dir> <control_dir> <output.ipk>

import os
import sys
import io
import tarfile

# 需要在 tar 中设为 755 权限的文件
EXEC_FILES = {
    './usr/bin/netflow',
    './usr/bin/netflow-rules',
    './usr/bin/netflow_x86_64',
    './usr/bin/netflow_aarch64',
    './usr/bin/netflow_arm',
    './usr/bin/netflow_mips',
    './usr/bin/netflow_mipsel',
    './etc/init.d/netflow',
    './etc/uci-defaults/luci-app-netflow',
}

# control.tar.gz 中的脚本全部需要 755
CONTROL_SCRIPTS = {'./preinst', './postinst', './prerm', './postrm'}


def fix_perms(tarinfo, is_control=False):
    """设置 Unix 权限"""
    if tarinfo.isdir():
        tarinfo.mode = 0o755
    elif is_control and tarinfo.name in CONTROL_SCRIPTS:
        tarinfo.mode = 0o755
    elif tarinfo.name in EXEC_FILES:
        tarinfo.mode = 0o755
    else:
        tarinfo.mode = 0o644
    return tarinfo


def make_tar_gz_bytes(source_dir, is_control=False):
    """将目录打包为 tar.gz，返回 bytes"""
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode='w:gz', format=tarfile.GNU_FORMAT) as tf:
        for root, dirs, files in os.walk(source_dir):
            for d in sorted(dirs):
                fpath = os.path.join(root, d)
                arcname = './' + os.path.relpath(fpath, source_dir).replace('\\', '/')
                ti = tf.gettarinfo(fpath, arcname)
                fix_perms(ti, is_control)
                ti.uid = ti.gid = 0
                ti.uname = ti.gname = 'root'
                tf.addfile(ti)

            for fn in sorted(files):
                fpath = os.path.join(root, fn)
                arcname = './' + os.path.relpath(fpath, source_dir).replace('\\', '/')
                ti = tf.gettarinfo(fpath, arcname)
                fix_perms(ti, is_control)
                ti.uid = ti.gid = 0
                ti.uname = ti.gname = 'root'
                with open(fpath, 'rb') as f:
                    tf.addfile(ti, f)
    return buf.getvalue()


def main():
    if len(sys.argv) < 4:
        print('用法: python repack_netflow_ipk.py <data_dir> <control_dir> <output.ipk>')
        sys.exit(1)

    data_dir = sys.argv[1]
    control_dir = sys.argv[2]
    output_ipk = sys.argv[3]

    # 1. 生成 debian-binary
    debian_data = b'2.0\n'

    # 2. 打包 control.tar.gz
    control_tgz = make_tar_gz_bytes(control_dir, is_control=True)

    # 3. 打包 data.tar.gz
    data_tgz = make_tar_gz_bytes(data_dir)

    # 4. 创建外层 tar.gz (IPK 格式)
    with tarfile.open(output_ipk, 'w:gz', format=tarfile.GNU_FORMAT) as outer:
        for name, data, mode in [
            ('./debian-binary', debian_data, 0o644),
            ('./control.tar.gz', control_tgz, 0o644),
            ('./data.tar.gz', data_tgz, 0o644),
        ]:
            ti = tarfile.TarInfo(name)
            ti.size = len(data)
            ti.mode = mode
            ti.uid = ti.gid = 0
            ti.uname = ti.gname = 'root'
            ti.mtime = 0
            outer.addfile(ti, io.BytesIO(data))

    size = os.path.getsize(output_ipk)
    print(f'IPK 打包完成: {output_ipk}')
    print(f'  大小: {size:,} bytes')
    print(f'  控制文件: {len(control_tgz):,} bytes')
    print(f'  数据文件: {len(data_tgz):,} bytes')


if __name__ == '__main__':
    main()