需求是:想要校验用户上传的压缩包中是否包含指定的文件夹和文件。
工具类目前支持:.zip .tar .tar.gz形式的压缩包文件
工具类如下:
pom中引入如下依赖:
<dependency>
<groupId>org.apache.commons</groupId>
<artifactId>commons-compress</artifactId>
<version>1.21</version> <!-- 根据需要选择合适的版本 -->
</dependency>
import org.apache.commons.compress.archivers.ArchiveEntry;
import org.apache.commons.compress.archivers.ArchiveException;
import org.apache.commons.compress.archivers.ArchiveInputStream;
import org.apache.commons.compress.archivers.ArchiveStreamFactory;
import org.apache.commons.compress.archivers.tar.TarArchiveEntry;
import org.apache.commons.compress.archivers.tar.TarArchiveInputStream;
import org.apache.commons.compress.archivers.zip.ZipArchiveEntry;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.HashSet;
import java.util.Set;
/**
* 用于检查镜像文件压缩包是否包含要求的文件夹,以及文件名
*/
public class CompressionChecker {
//此处,修改为自己需要检验的文件夹名称或者问价名称。
private static final String MODEL_FOLDER_NAME = "model/";
private static final String DOCKERFILE_NAME = "dockerfile";
private static final String INFERENCE_SERVICE_PY_NAME = "inference_services.py";
private static final String MODEL_DATA_IMPORT_PY_NAME = "model_data_import.py";
public static boolean checkCompression(String filePath) throws IOException {
File file = new File(filePath);
String fileName = file.getName().toLowerCase();
if (fileName.endsWith(".zip")) {
return checkZipFile(file);
} else if (fileName.endsWith(".tar") || fileName.endsWith(".tar.gz")) {
return checkTarFile(file);
}
// Unsupported file type
return false;
}
private static boolean checkZipFile(File file) throws IOException {
try (InputStream inputStream = new FileInputStream(file);
ArchiveInputStream archiveInputStream = new ArchiveStreamFactory().createArchiveInputStream(ArchiveStreamFactory.ZIP, inputStream)) {
Set<String> entries = new HashSet<>();
ZipArchiveEntry entry;
while ((entry = (ZipArchiveEntry) archiveInputStream.getNextEntry()) != null) {
processEntry(entry, archiveInputStream, entries);
}
return entries.contains(MODEL_FOLDER_NAME) && entries.contains(DOCKERFILE_NAME) &&
entries.contains(INFERENCE_SERVICE_PY_NAME) && entries.contains(MODEL_DATA_IMPORT_PY_NAME);
} catch (ArchiveException e) {
throw new RuntimeException(e);
}
}
private static boolean checkTarFile(File file) throws IOException {
try (InputStream inputStream = new FileInputStream(file)) {
Set<String> entries = new HashSet<>();
ArchiveInputStream tarInput;
if (file.getName().endsWith(".gz")) {
tarInput = new TarArchiveInputStream(new org.apache.commons.compress.compressors.gzip.GzipCompressorInputStream(inputStream));
} else {
tarInput = new TarArchiveInputStream(inputStream);
}
ArchiveEntry entry;
while ((entry = tarInput.getNextEntry()) != null) {
if (entry instanceof TarArchiveEntry) {
processEntry((TarArchiveEntry) entry, tarInput, entries);
}
}
return entries.contains(MODEL_FOLDER_NAME) && entries.contains(DOCKERFILE_NAME) &&
entries.contains(INFERENCE_SERVICE_PY_NAME) && entries.contains(MODEL_DATA_IMPORT_PY_NAME);
}
}
private static void processEntry(ZipArchiveEntry entry, ArchiveInputStream archiveInputStream, Set<String> entries) throws IOException {
if (entry.isDirectory()) {
String folderName = new File(entry.getName()).getName().toLowerCase();
entries.add(folderName+"/");
// 递归处理文件夹内部的内容
ArchiveEntry subEntry;
while ((subEntry = archiveInputStream.getNextEntry()) != null) {
if (subEntry.getName().startsWith(entry.getName())) {
processEntry((ZipArchiveEntry) subEntry, archiveInputStream, entries);
} else {
String fileName = new File(subEntry.getName()).getName().toLowerCase();
entries.add(fileName);
break;
}
}
} else {
String fileName = new File(entry.getName()).getName().toLowerCase();
entries.add(fileName);
}
}
private static void processEntry(TarArchiveEntry entry, ArchiveInputStream archiveInputStream, Set<String> entries) throws IOException {
if (entry.isDirectory()) {
String folderName = new File(entry.getName()).getName().toLowerCase();
entries.add(folderName+"/");
// 递归处理文件夹内部的内容
ArchiveEntry subEntry;
while ((subEntry = archiveInputStream.getNextEntry()) != null) {
if (subEntry.getName().startsWith(entry.getName())) {
processEntry((TarArchiveEntry) subEntry, archiveInputStream, entries);
} else {
String fileName = new File(subEntry.getName()).getName().toLowerCase();
entries.add(fileName);
break;
}
}
} else {
String fileName = new File(entry.getName()).getName().toLowerCase();
entries.add(fileName);
}
}
}