# 导入必要的模块
import os  # 用于操作系统路径
import sys  # 用于操作系统相关的功能
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/../../')  # 将父级目录添加到模块搜索路径中
import time  # 用于时间操作，例如延时

# 从client.portscan模块导入ShodanScan和NmapScan类，用于进行端口扫描
from client.portscan.ShodanScan import Scan
from client.portscan.NmapScan import Nmap_Portscan
# 从client.database模块导入数据库会话和数据表模型
from client.database import session, SrcAssets, SrcPorts

# 定义PortScan类，用于执行端口扫描操作
class PortScan:

    def __init__(self, ip):
        # 初始化方法，接收IP地址
        self.ip = ip

    def run(self):
        # 执行端口扫描
        # 1. 使用Shodan扫描IP的端口和漏洞
        port_list, vulns_list = Scan(ip=self.ip)
        # 2. 使用Nmap扫描获取更详细的端口信息
        port_dict = Nmap_Portscan(ip=self.ip, port_info_list=port_list)
        # 返回端口扫描结果和漏洞信息
        return port_dict, vulns_list

# 定义读取资产数据的函数
def ReadAssets():
    '''读取资产数据'''
    # 查询数据库，获取未进行端口扫描的资产（资产表中asset_port_flag为False）
    assets_sql = session.query(SrcAssets).filter(SrcAssets.asset_port_flag == False).first()
    # 提交数据库事务
    session.commit()

    if assets_sql:
        # 获取该资产的IP地址
        ip = assets_sql.asset_ip
        # 查询所有具有相同IP地址的资产
        assets_sql1 = session.query(SrcAssets).filter(SrcAssets.asset_ip == ip).all()

        # 遍历所有相同IP的资产，将其端口扫描状态更新为True
        for sql in assets_sql1:
            sql.asset_port_flag = True  # 修改资产的端口扫描状态为True
            session.add(sql)  # 将修改的资产添加到会话中

        try:
            # 提交事务，将修改保存到数据库
            session.commit()
        except Exception as error:
            # 如果提交失败，打印异常信息并回滚事务
            print(f'[-]端口扫描-修改IP扫描状态异常{error}')
            session.rollback()

    # 返回读取到的资产数据
    return assets_sql

# 定义将扫描结果写入数据库的函数
def WritePosts(port_dict, assets_sql):
    '''端口扫描入库'''
    # 遍历扫描到的端口信息
    for info in port_dict:
        # 构造要插入到数据库的端口数据
        port_sql = SrcPorts(
            port_name=assets_sql.asset_name,  # 资产名称
            port_host=assets_sql.asset_host,  # 资产主机
            port_ip=assets_sql.asset_ip,  # 资产IP地址
            port_port=port_dict[info]['port'],  # 端口号
            port_service=port_dict[info]['name'],  # 服务名称
            port_product=port_dict[info]['product'],  # 产品名称
            port_version=port_dict[info]['version']  # 产品版本
        )
        # 将端口信息添加到数据库会话中
        session.add(port_sql)
        try:
            # 提交事务，将端口信息保存到数据库
            session.commit()
        except Exception as error:
            # 如果提交失败，回滚事务并打印错误信息
            session.rollback()
            print(f'[-]端口入库异常{error}')
    
    # 打印扫描结果已入库完成的提示
    print(f'[+]端口[{assets_sql.asset_ip}]入库完成')

# 主函数，启动端口扫描和入库操作
def main():
    print('[+]端口扫描启动')
    while True:
        # 读取资产数据
        assets_sql = ReadAssets()

        # 如果没有可处理的资产，等待30秒后继续尝试
        if not assets_sql:
            time.sleep(30)
        else:
            # 创建PortScan实例，传入要扫描的IP地址
            portscan = PortScan(assets_sql.asset_ip)
            # 运行端口扫描，获取端口信息和漏洞信息
            port_dict, vulns_list = portscan.run()

            # 如果扫描结果不为空，写入数据库
            if port_dict:
                WritePosts(port_dict, assets_sql)

# 如果该脚本作为主程序执行，启动端口扫描程序
if __name__ == '__main__':
    main()