用腾讯云AI代码助手开发一款数据库敏感信息检查工具

引言

作为一个优秀的 IT 技术民工,需要始终学习先进的技术并将技术转换为生产力,目前 AI 领域编码辅助工具层出不穷,开发者有必要或者说必须要掌握相关工具的使用,以提高编码效率,降低编码错误。这次我通过一个简单的项目和大家分享一下腾讯云 AI 代码助手的使用。

AI 开发实践

开发环境介绍

目前我使用的开发语言是Java,主要在 IntelliJ IDEA 开发,选用 Java 开发的原因主要是个人习惯,对 Java 比较熟悉,不过也是第一次去开发类似腾讯这种数据库扫描工具,类似 T- Sec 数据安全审计,挖掘数据库运行过程中各类潜在风险和隐患的扫描小工具,还是有些挑战。所以我选择采用目前对比体验效果比较不错的腾讯云AI代码助手

首先我们从安装插件开始,打开 IntelliJ IDEA ,通过 Settings->Plugins->搜索腾讯云 AI 代码助手,点击 Install

 

然后,右下角点击 icon 微信扫码登录

开发背景简要说明

开发数据库扫描工具的原因是因为我需要对公司内部所有运行在腾讯云上面的数据库进行检查。

信息安全:

要求不允许数据库中存放任何敏感信息,包括身份证号,手机号,银行卡号,邮件地址等信息

面临挑战:

由于数据库众多,数据量巨大,如果进行人工核验非常耗时耗力,最重要的是难免会有所遗漏

解决方案:

基于信息安全和挑战,决定采用 AI 来开发一个简单的工具来实现此功能,提升自己工作效率,数据库铭感信息检查工具主要包括以下功能:

  • 面向多个腾讯云数据库连接,工具可以一次性接入并完成所有数据库的扫描,以提升检查效率
  • 可以自动获取一个腾讯云数据库下的多个 DB,并进行数据扫描
  • 能匹配数据格式是否是敏感信息
  • 将识别到的敏感信息输出日志,以便后续反馈给相关人员进行处理

编码实现

首先定义多个数据库的连接方式

以下为 yaml demo 示例文件格式,这种定义方式可以将多个数据库信息写在一起,一次扫描所有数据库服务器

databases:
  - host: 192.168.1.2
    port: 3306
    username: root
    password: MysqlPasswd
  - host: 192.168.1.3
    port: 3306
    username: root
    password: MysqlPasswd

然后,进行定义 class 类用于绑定 yaml,将 yaml 文件的数据库信息映射成 DatabaseInfo 对象,方便我后续的操作

package org.checkdb;

public class DatabaseInfo {
    String databaseHost;
    String databasePort;
    String databaseUser;
    String databasePassword;

    // Getter and Setter methods
    public String getHost() {
        return databaseHost;
    }

    public void setHost(String databaseHost) {
        this.databaseHost = databaseHost;
    }

    public String getPort() {
        return databasePort;
    }

    public void setPort(String databasePort) {
        this.databasePort = databasePort;
    }

    public String getUsername() {
        return databaseUser;
    }

    public void setUsername(String databaseUser) {
        this.databaseUser = databaseUser;
    }

    public String getPassword() {
        return databasePassword;
    }

    public void setPassword(String databasePassword) {
        this.databasePassword = databasePassword;
    }
}

在 mian 中绑定 yaml 和类,在工具启动时映射 dataeaseinfo 对象中,这部分代码我们使用腾讯云 AI 代码助手帮我们生成一个简单的例子,我们简单修改完成。我们可以通过问答的方式获取相关例子

通过腾讯云 AI 代码助手 AI 技术问答给出的例子,基本简单修改即可完成使用,但是注意别忘了添加包依赖

package org.checkdb.utils;

import org.checkdb.DatabaseInfo;
import org.yaml.snakeyaml.Yaml;

import java.io.InputStream;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

public class DatabaseMapper {
    private static final Yaml yaml = new Yaml();

    public static List<DatabaseInfo> mapDatabasesFromYaml(String yamlFilePath) {
        InputStream inputStream = DatabaseMapper.class.getClassLoader().getResourceAsStream(yamlFilePath);
        Map<String, List<Map<String, String>>> databasesMap = yaml.load(inputStream);
        List<DatabaseInfo> databaseInfos = new ArrayList<>();

        for (Map<String, String> databaseMap : databasesMap.get("databases")) {
            DatabaseInfo databaseInfo = new DatabaseInfo();
            databaseInfo.setHost(databaseMap.get("host"));
            databaseInfo.setPort(String.valueOf(databaseMap.get("port")));
            databaseInfo.setUsername(databaseMap.get("username"));
            databaseInfo.setPassword(databaseMap.get("password"));
            databaseInfos.add(databaseInfo);
        }

        return databaseInfos;
    }
}

在 main 方法中增加调用,这样我们就在 mian 方法中获取到了所有的数据库信息 List

   public static void main(String[] args) {
        List<DatabaseInfo> databaseInfos = DatabaseMapper.mapDatabasesFromYaml("databases.yaml");
        for (DatabaseInfo databaseInfo : databaseInfos) {
            List<String> databases = DatabaseConnector.getAllDatabases(databaseInfo);
        }
    }

编写连接数据库的方法

获取单个 mysql 中的所有 database,然后我们把这个方法放到 main 中的,对 databaseInfos 进行循环,依次获取所有 mysql 数据库中的所有 database。

现在继续使用腾讯云 AI 代码助手帮助我们提供示例

    public static List<String> getAllDatabases(DatabaseInfo databaseInfo) {
        List<String> databases = new ArrayList<>();
        try {
            Connection conn = DriverManager.getConnection("jdbc:mysql://" + databaseInfo.getHost() + ":" + databaseInfo.getPort()+ "?useSSL=false&serverTimezone=UTC", databaseInfo.getUsername(), databaseInfo.getPassword());
            Statement stmt = conn.createStatement();
            ResultSet rs = stmt.executeQuery("SHOW DATABASES");
            while (rs.next()) {
                databases.add(rs.getString(1));
            }
            conn.close();
        } catch (SQLException e) {
            e.printStackTrace();
        }
        return databases;
    }

在 main 方法中增加调用,现在我们可以依次获取所有 mysql 数据库中的所有 database

   public static void main(String[] args) {
        List<DatabaseInfo> databaseInfos = DatabaseMapper.mapDatabasesFromYaml("databases.yaml");
        for (DatabaseInfo databaseInfo : databaseInfos) {
            List<String> databases = DatabaseConnector.getAllDatabases(databaseInfo);
        }

获取单个 database 中的所有表,这次我们换一种新的方式使用腾讯云 AI 代码助手,通过编写代码注释,让腾讯云 AI 代码助手帮我们生成代码

    public static List<String> getAllTables(DatabaseInfo databaseInfo, String databaseName) {
        List<String> tables = new ArrayList<>();
        try {
            Connection conn = DriverManager.getConnection("jdbc:mysql://" + databaseInfo.getHost() + ":" + databaseInfo.getPort()+ "?useSSL=false&serverTimezone=UTC", databaseInfo.getUsername(), databaseInfo.getPassword());
            Statement stmt = conn.createStatement();
            ResultSet rs = stmt.executeQuery("SHOW TABLES FROM " + databaseName);
            while (rs.next()) {
                tables.add(rs.getString(1));
            }
            conn.close();
        } catch (SQLException e) {
            e.printStackTrace();
        }
        return tables;
    }

在 main 方法中增加调用,现在我们可以依次获取所有 mysql 数据库中的所有 database 的所有 table

    public static void main(String[] args) {
        List<DatabaseInfo> databaseInfos = DatabaseMapper.mapDatabasesFromYaml("databases.yaml");
        for (DatabaseInfo databaseInfo : databaseInfos) {
            List<String> databases = DatabaseConnector.getAllDatabases(databaseInfo);
            for (String database : databases) {
                List<String> tables = DatabaseConnector.getAllTables(databaseInfo,database);
            }
        }
    }

数据格式校验方法

现在已经有了所有 database以及 table,我继续编写效验身份证格式,电话号格式和银行卡格式的方法,也继续使用通过注释的方式生成代码

以下为 AI 生成的代码,达到了 100% 可用

package org.checkdb.utils;

public class CheckRE {

    // 判断数据格式是否是身份证号
    public static boolean isIdCard(String idCard) {
        return idCard.matches("^[1-9]\\d{5}(18|19|([23]\\d))\\d{2}((0[1-9])|(10|11|12))(([0-2][1-9])|10|20|30|31)\\d{3}[0-9Xx]$");
    }

    // 判断数据格式是否是手机号
    public static boolean isMobileNO(String mobiles) {
        return mobiles.matches("^((13[0-9])|(14[5,7])|(15[0-3,5-9])|(17[0,3,5-8])|(18[0-9])|166|198|199|(147))\\d{8}$");
    }

    // 判断数据格式是否是银行卡号
    public static boolean isBankCard(String bankCard) {
        return bankCard.matches("^[1-9]\\d{9,29}$");
    }
}

注:这些判断方法肯定不是非常严格的验证,但是用于判断数据库中是否存在这些敏感类型的数据是足够了

扫描逻辑

最后,通过编写工具扫描逻辑,由于数据表中的数据非常多,全表扫描存在性能问题,因此并不可取,本次,我们以每次获取每张表的前 500 条数据进行判断,如果存在敏感数据,我们记录到文件中,依然腾讯云 AI 代码助手生成代码

package org.checkdb.utils;

import org.checkdb.DatabaseInfo;

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.sql.*;
import java.util.ArrayList;
import java.util.Objects;

import static org.checkdb.utils.CheckRE.*;

public class ScanDatabase {

    public static void ScanTable(DatabaseInfo databaseInfo,String databaseName,String tableName){

        if (Objects.equals(databaseName, "mysql") || Objects.equals(databaseName, "sys") || Objects.equals(databaseName, "performance_schema") || Objects.equals(databaseName, "information_schema")) {
            return;
        }
        Connection conn = null;
        Statement stmt = null;
        ResultSet rs = null;
        try {
            // 连接数据库
            conn = DriverManager.getConnection("jdbc:mysql://" + databaseInfo.getHost() + ":" + databaseInfo.getPort()+"/"+databaseName+ "?useSSL=false&serverTimezone=UTC", databaseInfo.getUsername(), databaseInfo.getPassword());
            // 获取表的所有字段信息
            DatabaseMetaData metaData = conn.getMetaData();
            ResultSet fieldsRS = metaData.getColumns(null, null, tableName, null);
            // 存储字段名
            ArrayList<String> fields = new ArrayList<>();
            while (fieldsRS.next()) {
                fields.add(fieldsRS.getString("COLUMN_NAME"));
            }
            // 执行SQL查询,获取前500条数据
            String sql = "SELECT * FROM " + tableName + " LIMIT 500";
            stmt = conn.createStatement();
            rs = stmt.executeQuery(sql);
            // 判断每个字段是否包含敏感信息并写入文件
            BufferedWriter writer = new BufferedWriter(new FileWriter("sensitive_data.txt", true));
            while (rs.next()) {
                for (String field : fields) {
                    String value = rs.getString(field);
                    if (value == null) {
                        continue;
                    }
                    if (isIdCard(value)) {
                        writer.write("数据库IP: "+ databaseInfo.getHost() + "数据库名称: "+ databaseName + "表名称:" + tableName + "字段名: " + field + "存在敏感数据: " + value + "\n");
                        break;
                    }
                    if (isMobileNO(value)) {
                        writer.write("数据库IP: "+ databaseInfo.getHost() + "数据库名称: "+ databaseName + "表名称:" + tableName + "字段名: " + field + "存在敏感数据: " + value + "\n");
                        break;
                    }
                    if (isBankCard(value)) {
                        writer.write("数据库IP: "+ databaseInfo.getHost() + "数据库名称: "+ databaseName + "表名称:" + tableName + "字段名: " + field + "存在敏感数据: " + value + "\n");
                        break;
                    }
                }
            }
            writer.close();
        } catch (SQLException | IOException e) {
            e.printStackTrace();
        } finally {
            // 关闭资源
            try {
                if (rs != null) rs.close();
                if (stmt != null) stmt.close();
                if (conn != null) conn.close();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }
}
             

在 main 方法中增加调用

public class Main {
    public static void main(String[] args) {
        List<DatabaseInfo> databaseInfos = DatabaseMapper.mapDatabasesFromYaml("databases.yaml");
        for (DatabaseInfo databaseInfo : databaseInfos) {
            List<String> databases = DatabaseConnector.getAllDatabases(databaseInfo);
            for (String database : databases) {
                List<String> tables = DatabaseConnector.getAllTables(databaseInfo,database);
                for (String table : tables){
                    ScanDatabase.ScanTable(databaseInfo,database,table);
                }
            }
        }
    }
}

注:1、要排除 Mysql 自带库;2、要考虑查询值空指针

 Code Review

现在我们代码逻辑和主要功能都已经实现,但是还是需要进行 code review,以便检查出代码中存在的潜在问题,还好腾讯云 AI 代码助手提供了代码优化功能,我们使用一下看看效果,此处主要的两个功能是,代码优化和缺陷检查

通过代码优化,对 scanDatabase 进行优化

该功能给出了一下优化建议,以及示例代码,我们可以参考和建议进行修改

  • 使用 try-with-resources 语句自动关闭资源,避免潜在的资源泄露。
  • 将敏感信息的检测逻辑封装成方法,提高代码的可读性和可维护性。
  • 使用预编译语句(PreparedStatement)来执行 SQL 查询,提高代码的安全性和性能。
  • 避免在循环中频繁地打开和关闭文件,可以考虑使用 try-with-resources 语句来管理文件资源。

通过缺陷检查,我们对 scanDatabase 进行缺陷检查

该功能给出的建议如下

SQL 注入风险:原代码中直接拼接 SQL 查询语句,存在 SQL 注入的风险。为了避免这个问题,应该使用 PreparedStatement 来代替 Statement。 

资源关闭异常:在 finally 块中关闭资源时,如果任何一个资源关闭失败,后续的资源关闭操作将不会被执行。应该分别捕获每个资源的关闭异常。 

文件写入异常处理:在写入文件时,如果发生异常,可能会导致文件没有被正确关闭。应该使用 try-with-resources 语句来确保文件在异常发生时也能被正确关闭。 

数据库连接字符串硬编码:数据库连接字符串中的 useSSL = false&serverTimezone = UTC 可能不适用于所有情况,应该允许通过参数传递。 

文件写入路径硬编码:写入文件的路径被硬编码为 sensitive_data.txt,这可能导致文件被覆盖或写入到不期望的位置。应该允许通过参数传递文件路径。

Code Review小结

通过使用代码优化、代码缺陷检查,可以帮助我快速发现解决问题,也是对日常编码过程的一些宝贵建议,让我的思路可以更开阔,但 AI 给出的建议并非是一定要采纳的,需要根据情况具体分析,比如存在 sql 注入风险,所有 sql 均在代码内部,并不接收任何参数作为 sql 的一部分,则可暂不考虑优化。

工具效果

通过创建了两个 MySQL 数据库,创造一些促五示例数据,并进行扫描检测,可看到工具已查询出存在敏感字段的内容

注:为了验证和企业内部数据的安全性,这里通过在一台服务器上部署创建两个测试的数据库服务器来验证。

获得帮助与提升

通过使用腾讯云 AI 代码助手的问答、代码补全、代码优化、缺陷检查等功能,我们顺利完成了这个小工具的开发,只能说功能非常好用,极大的提高 Coding 效率,基于代码上下文逻辑,腾讯云 AI 代码助手可以自动补全,通过问答方式也可快速获取想要的例子。

本次开发通过 AI 技术问答的方式解决了一部分代码逻辑设计的问题(因为我没有思考如何设计,只是把问题抛给腾讯云 AI 代码助手)可以说是小白式开发;还有很多的功能我实际用到了,但是文章中没有体现出来,比如下图出现报错时,腾讯云 AI 代码助手可以帮助定位问题,另外,就是腾讯云 AI 代码助手可以快速帮你生成代码注释,简直懒人神器,最后祝愿腾讯云 AI 代码助手越来越棒!

posted @   粤海科技君  阅读(36)  评论(0编辑  收藏  举报
相关博文:
阅读排行:
· CSnakes vs Python.NET:高效嵌入与灵活互通的跨语言方案对比
· 【.NET】调用本地 Deepseek 模型
· Plotly.NET 一个为 .NET 打造的强大开源交互式图表库
· 上周热点回顾(2.17-2.23)
· 如何使用 Uni-app 实现视频聊天(源码,支持安卓、iOS)
点击右上角即可分享
微信分享提示