引言

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

AI 开发实践

开发环境介绍

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

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

开发背景简要说明

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

信息安全：

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

面临挑战：

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

解决方案：

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

面向多个数据库连接，工具可以一次性接入并完成所有数据库的扫描，以提升检查效率

可以自动获取一个数据库下的多个 DB，并进行数据扫描

能匹配数据格式是否是敏感信息

将识别到的敏感信息输出日志，以便后续反馈给相关人员进行处理

编码实现

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

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

databases:  -host:192.168.1.2port:3306username: rootpassword:MysqlPasswd-host:192.168.1.3port:3306username: rootpassword:MysqlPasswd

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

package org.checkdb;publicclassDatabaseInfo{StringdatabaseHost;StringdatabasePort;StringdatabaseUser;StringdatabasePassword;// Getter and Setter methodspublicStringgetHost() {returndatabaseHost;    }    publicvoidsetHost(StringdatabaseHost) {this.databaseHost= databaseHost;    }    publicStringgetPort() {returndatabasePort;    }    publicvoidsetPort(StringdatabasePort) {this.databasePort= databasePort;    }    publicStringgetUsername() {returndatabaseUser;    }    publicvoidsetUsername(StringdatabaseUser) {this.databaseUser= databaseUser;    }    publicStringgetPassword() {returndatabasePassword;    }    publicvoidsetPassword(StringdatabasePassword) {this.databasePassword= databasePassword;    }}

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

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

packageorg.checkdb.utils;importorg.checkdb.DatabaseInfo;importorg.yaml.snakeyaml.Yaml;importjava.io.InputStream;importjava.util.ArrayList;importjava.util.List;importjava.util.Map;publicclassDatabaseMapper{    privatestaticfinalYamlyaml =newYaml();    publicstaticListmapDatabasesFromYaml(StringyamlFilePath) {InputStreaminputStream =DatabaseMapper.class.getClassLoader().getResourceAsStream(yamlFilePath);Map>> databasesMap = yaml.load(inputStream);List databaseInfos =newArrayList<>();for(Map databaseMap : databasesMap.get("databases")) {DatabaseInfodatabaseInfo =newDatabaseInfo();            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);        }returndatabaseInfos;    }}

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

publicstaticvoidmain(String[] args) {List databaseInfos =DatabaseMapper.mapDatabasesFromYaml("databases.yaml");for(DatabaseInfodatabaseInfo : databaseInfos) {List databases =DatabaseConnector.getAllDatabases(databaseInfo);        }    }

编写连接数据库的方法

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

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

publicstaticListgetAllDatabases(DatabaseInfo databaseInfo) {List databases =newArrayList<>();try{Connectionconn =DriverManager.getConnection("jdbc:mysql://"+ databaseInfo.getHost() +":"+ databaseInfo.getPort()+"?useSSL=false&serverTimezone=UTC", databaseInfo.getUsername(), databaseInfo.getPassword());Statementstmt = conn.createStatement();ResultSetrs = stmt.executeQuery("SHOW DATABASES");while(rs.next()) {                databases.add(rs.getString(1));            }            conn.close();        }catch(SQLExceptione) {            e.printStackTrace();        }returndatabases;    }

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

publicstaticvoidmain(String[] args) {List databaseInfos =DatabaseMapper.mapDatabasesFromYaml("databases.yaml");for(DatabaseInfodatabaseInfo : databaseInfos) {List databases =DatabaseConnector.getAllDatabases(databaseInfo);        }

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

publicstaticListgetAllTables(DatabaseInfo databaseInfo,StringdatabaseName) {List tables =newArrayList<>();try{Connectionconn =DriverManager.getConnection("jdbc:mysql://"+ databaseInfo.getHost() +":"+ databaseInfo.getPort()+"?useSSL=false&serverTimezone=UTC", databaseInfo.getUsername(), databaseInfo.getPassword());Statementstmt = conn.createStatement();ResultSetrs = stmt.executeQuery("SHOW TABLES FROM "+ databaseName);while(rs.next()) {                tables.add(rs.getString(1));            }            conn.close();        }catch(SQLExceptione) {            e.printStackTrace();        }returntables;    }

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

publicstaticvoidmain(String[] args) {List databaseInfos =DatabaseMapper.mapDatabasesFromYaml("databases.yaml");for(DatabaseInfodatabaseInfo : databaseInfos) {List databases =DatabaseConnector.getAllDatabases(databaseInfo);for(Stringdatabase : databases) {List tables =DatabaseConnector.getAllTables(databaseInfo,database);            }        }    }

数据格式校验方法

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

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

package org.checkdb.utils;publicclassCheckRE{// 判断数据格式是否是身份证号publicstaticbooleanisIdCard(StringidCard) {returnidCard.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]$");    }// 判断数据格式是否是手机号publicstaticbooleanisMobileNO(Stringmobiles) {returnmobiles.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}$");    }// 判断数据格式是否是银行卡号publicstaticbooleanisBankCard(StringbankCard) {returnbankCard.matches("^[1-9]\\d{9,29}$");    }}

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

扫描逻辑

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

packageorg.checkdb.utils;importorg.checkdb.DatabaseInfo;importjava.io.BufferedWriter;importjava.io.FileWriter;importjava.io.IOException;importjava.sql.*;importjava.util.ArrayList;importjava.util.Objects;importstaticorg.checkdb.utils.CheckRE.*;publicclassScanDatabase{    publicstaticvoidScanTable(DatabaseInfodatabaseInfo,StringdatabaseName,StringtableName){if(Objects.equals(databaseName,"mysql") ||Objects.equals(databaseName,"sys") ||Objects.equals(databaseName,"performance_schema") ||Objects.equals(databaseName,"information_schema")) {return;        }Connectionconn =null;Statementstmt =null;ResultSetrs =null;try{// 连接数据库conn =DriverManager.getConnection("jdbc:mysql://"+ databaseInfo.getHost() +":"+ databaseInfo.getPort()+"/"+databaseName+"?useSSL=false&serverTimezone=UTC", databaseInfo.getUsername(), databaseInfo.getPassword());// 获取表的所有字段信息DatabaseMetaDatametaData = conn.getMetaData();ResultSetfieldsRS = metaData.getColumns(null,null, tableName,null);// 存储字段名ArrayList fields =newArrayList<>();while(fieldsRS.next()) {                fields.add(fieldsRS.getString("COLUMN_NAME"));            }// 执行SQL查询，获取前500条数据Stringsql ="SELECT * FROM "+ tableName +" LIMIT 500";            stmt = conn.createStatement();            rs = stmt.executeQuery(sql);// 判断每个字段是否包含敏感信息并写入文件BufferedWriterwriter =newBufferedWriter(newFileWriter("sensitive_data.txt",true));while(rs.next()) {for(Stringfield : fields) {Stringvalue = 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|IOExceptione) {            e.printStackTrace();        }finally{// 关闭资源try{if(rs !=null) rs.close();if(stmt !=null) stmt.close();if(conn !=null) conn.close();            }catch(SQLExceptione) {                e.printStackTrace();            }        }    }}

在 main 方法中增加调用

publicclassMain{    publicstaticvoidmain(String[] args) {List databaseInfos =DatabaseMapper.mapDatabasesFromYaml("databases.yaml");for(DatabaseInfodatabaseInfo : databaseInfos) {List databases =DatabaseConnector.getAllDatabases(databaseInfo);for(Stringdatabase : databases) {List tables =DatabaseConnector.getAllTables(databaseInfo,database);for(Stringtable : 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 数据库，创造一些促五示例数据，并进行扫描检测，可看到工具已查询出存在敏感字段的内容

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