用腾讯云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)