From ad78dad7c708e9cb110cc418bf9e5de4b91bf133 Mon Sep 17 00:00:00 2001 From: chn Date: Fri, 3 May 2024 16:15:40 +0800 Subject: [PATCH] =?UTF-8?q?=E5=9F=BA=E6=9C=AC=E5=AE=8C=E6=88=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- CMakeLists.txt | 9 +- flake.nix | 6 +- include/hpcstat/common.hpp | 15 ++ include/hpcstat/env.hpp | 12 ++ include/hpcstat/keys.hpp | 10 ++ include/hpcstat/lfs.hpp | 14 ++ include/hpcstat/sql.hpp | 87 ++++++++++++ include/hpcstat/ssh.hpp | 13 ++ keys/00 | 1 + keys/01 | 1 + keys/02 | 1 + keys/03 | 1 + keys/chn | 1 + keys/gb | 1 + keys/xll | 1 + keys/yjq | 1 + keys/zem | 1 + src/common.cpp | 35 +++++ src/env.cpp | 19 +++ src/keys.cpp | 17 +++ src/lfs.cpp | 100 ++++++++++++++ src/main.cpp | 271 ++++++++++++------------------------- src/sql.cpp | 55 ++++++++ src/ssh.cpp | 84 ++++++++++++ 24 files changed, 566 insertions(+), 190 deletions(-) create mode 100644 include/hpcstat/common.hpp create mode 100644 include/hpcstat/env.hpp create mode 100644 include/hpcstat/keys.hpp create mode 100644 include/hpcstat/lfs.hpp create mode 100644 include/hpcstat/sql.hpp create mode 100644 include/hpcstat/ssh.hpp create mode 100644 keys/00 create mode 100644 keys/01 create mode 100644 keys/02 create mode 100644 keys/03 create mode 100644 keys/chn create mode 100644 keys/gb create mode 100644 keys/xll create mode 100644 keys/yjq create mode 100644 keys/zem create mode 100644 src/common.cpp create mode 100644 src/env.cpp create mode 100644 src/keys.cpp create mode 100644 src/lfs.cpp create mode 100644 src/sql.cpp create mode 100644 src/ssh.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 4994c52..94b485c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -12,10 +12,15 @@ endif() find_package(fmt REQUIRED) find_package(Boost REQUIRED COMPONENTS headers filesystem) find_package(zxorm REQUIRED) +find_package(nlohmann_json REQUIRED) +find_path(ZPP_BITS_INCLUDE_DIR zpp_bits.h REQUIRED) +find_package(range-v3 REQUIRED) -add_executable(hpcstat src/main.cpp) +add_executable(hpcstat src/main.cpp src/env.cpp src/keys.cpp src/ssh.cpp src/sql.cpp src/lfs.cpp src/common.cpp) target_compile_features(hpcstat PUBLIC cxx_std_23) -target_link_libraries(hpcstat PRIVATE fmt::fmt Boost::headers Boost::filesystem zxorm::zxorm) +target_include_directories(hpcstat PRIVATE ${PROJECT_SOURCE_DIR}/include ${ZPP_BITS_INCLUDE_DIR}) +target_link_libraries(hpcstat PRIVATE fmt::fmt Boost::headers Boost::filesystem zxorm::zxorm + nlohmann_json::nlohmann_json range-v3::range-v3) install(TARGETS hpcstat RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}) diff --git a/flake.nix b/flake.nix index 05d293f..cd7c192 100644 --- a/flake.nix +++ b/flake.nix @@ -22,7 +22,8 @@ { name = "hpcstat"; src = ./.; - buildInputs = with pkgs.pkgsStatic; [ boost fmt localPackages.zxorm ]; + buildInputs = with pkgs.pkgsStatic; + [ boost fmt localPackages.zxorm nlohmann_json localPackages.zpp-bits range-v3 ]; nativeBuildInputs = with pkgs; [ cmake pkg-config ]; postInstall = "cp ${openssh}/bin/ssh-add $out/bin"; }; @@ -32,7 +33,8 @@ devShell.x86_64-linux = pkgs.mkShell { nativeBuildInputs = with pkgs; [ pkg-config cmake clang-tools_18 ]; - buildInputs = (with pkgs.pkgsStatic; [ fmt boost localPackages.zxorm ]); + buildInputs = (with pkgs.pkgsStatic; + [ fmt boost localPackages.zxorm nlohmann_json localPackages.zpp-bits range-v3 ]); # hardeningDisable = [ "all" ]; # NIX_DEBUG = "1"; CMAKE_EXPORT_COMPILE_COMMANDS = "1"; diff --git a/include/hpcstat/common.hpp b/include/hpcstat/common.hpp new file mode 100644 index 0000000..7c9b4bf --- /dev/null +++ b/include/hpcstat/common.hpp @@ -0,0 +1,15 @@ +# pragma once +# include +# include +# include +# include + +namespace hpcstat +{ + // run a program, wait until it exit, return its stdout if it return 0, otherwise nullopt + std::optional exec + (std::filesystem::path program, std::vector args, std::optional stdin = std::nullopt); + + // get current time + long now(); +} diff --git a/include/hpcstat/env.hpp b/include/hpcstat/env.hpp new file mode 100644 index 0000000..015232c --- /dev/null +++ b/include/hpcstat/env.hpp @@ -0,0 +1,12 @@ +# pragma once +# include +# include + +namespace hpcstat::env +{ + // check if the program is running in an interactive shell + bool interactive(); + + // get the value of an environment variable + std::optional env(std::string name, bool required = false); +} diff --git a/include/hpcstat/keys.hpp b/include/hpcstat/keys.hpp new file mode 100644 index 0000000..56945c7 --- /dev/null +++ b/include/hpcstat/keys.hpp @@ -0,0 +1,10 @@ +# pragma once +# include +# include + +namespace hpcstat +{ + // valid keys + struct Key { std::string PubkeyFilename; std::string Username; }; + extern std::map Keys; +} diff --git a/include/hpcstat/lfs.hpp b/include/hpcstat/lfs.hpp new file mode 100644 index 0000000..587a940 --- /dev/null +++ b/include/hpcstat/lfs.hpp @@ -0,0 +1,14 @@ +# pragma once +# include +# include +# include +# include +# include + +namespace hpcstat::lfs +{ + std::optional> bsub(std::vector args); + // JobId -> { SubmitTime, Status, CpuTime } + std::optional>> bjobs_list(); + std::optional bjobs_detail(unsigned jobid); +} diff --git a/include/hpcstat/sql.hpp b/include/hpcstat/sql.hpp new file mode 100644 index 0000000..1e3db30 --- /dev/null +++ b/include/hpcstat/sql.hpp @@ -0,0 +1,87 @@ +# pragma once +# include +# include +# include + +namespace hpcstat::sql +{ + struct LoginData + { + unsigned Id = 0; long Time; + std::string Key, SessionId, Signature = ""; + std::optional Subaccount, Ip; + bool Interactive; + using serialize = zpp::bits::members<8>; + }; + using LoginTable = zxorm::Table + < + "login", LoginData, + zxorm::Column<"id", &LoginData::Id, zxorm::PrimaryKey<>>, + zxorm::Column<"time", &LoginData::Time>, + zxorm::Column<"key", &LoginData::Key>, + zxorm::Column<"session_id", &LoginData::SessionId>, + zxorm::Column<"signature", &LoginData::Signature>, + zxorm::Column<"sub_account", &LoginData::Subaccount>, + zxorm::Column<"ip", &LoginData::Ip>, + zxorm::Column<"interactive", &LoginData::Interactive> + >; + struct LogoutData { unsigned Id = 0; long Time; std::string SessionId; }; + using LogoutTable = zxorm::Table + < + "logout", LogoutData, + zxorm::Column<"id", &LogoutData::Id, zxorm::PrimaryKey<>>, + zxorm::Column<"time", &LogoutData::Time>, + zxorm::Column<"sessionid", &LogoutData::SessionId> + >; + struct SubmitJobData + { + unsigned Id = 0; + long Time; + unsigned JobId; + std::string Key, SessionId, SubmitDir, JobCommand, Signature = ""; + std::optional Subaccount, Ip; + using serialize = zpp::bits::members<10>; + }; + using SubmitJobTable = zxorm::Table + < + "submitjob", SubmitJobData, + zxorm::Column<"id", &SubmitJobData::Id, zxorm::PrimaryKey<>>, + zxorm::Column<"time", &SubmitJobData::Time>, + zxorm::Column<"job_id", &SubmitJobData::JobId>, + zxorm::Column<"key", &SubmitJobData::Key>, + zxorm::Column<"session_id", &SubmitJobData::SessionId>, + zxorm::Column<"submit_dir", &SubmitJobData::SubmitDir>, + zxorm::Column<"job_command", &SubmitJobData::JobCommand>, + zxorm::Column<"signature", &SubmitJobData::Signature>, + zxorm::Column<"sub_account", &SubmitJobData::Subaccount>, + zxorm::Column<"ip", &SubmitJobData::Ip> + >; + struct FinishJobData + { + unsigned Id = 0; + long Time; + unsigned JobId; + std::string JobResult, SubmitTime, Signature = ""; + double CpuTime; + using serialize = zpp::bits::members<7>; + }; + using FinishJobTable = zxorm::Table + < + "finishjob", FinishJobData, + zxorm::Column<"id", &FinishJobData::Id, zxorm::PrimaryKey<>>, + zxorm::Column<"time", &FinishJobData::Time>, + zxorm::Column<"job_id", &FinishJobData::JobId>, + zxorm::Column<"job_result", &FinishJobData::JobResult>, + zxorm::Column<"submit_time", &FinishJobData::SubmitTime>, + zxorm::Column<"signature", &FinishJobData::Signature>, + zxorm::Column<"cpu_time", &FinishJobData::CpuTime> + >; + // 序列化任意数据，用于之后签名 + std::string serialize(auto data); + // 初始化数据库 + bool initdb(); + // 将数据写入数据库 + bool writedb(auto value); + // 查询 bjobs -a 的结果中，有哪些是已经被写入到数据库中的（按照任务 id 和提交时间计算），返回未被写入的任务 id + std::optional> finishjob_remove_existed(std::map jobid_submit_time); +} diff --git a/include/hpcstat/ssh.hpp b/include/hpcstat/ssh.hpp new file mode 100644 index 0000000..0a033fa --- /dev/null +++ b/include/hpcstat/ssh.hpp @@ -0,0 +1,13 @@ +# pragma once +# include +# include + +namespace hpcstat::ssh +{ + // get a valid public key fingerprint + std::optional fingerprint(); + // sign a message with the key of specified fingerprint + std::optional sign(std::string message, std::string fingerprint); + // verify a message with the key of specified fingerprint + bool verify(std::string message, std::string signature, std::string fingerprint); +} diff --git a/keys/00 b/keys/00 new file mode 100644 index 0000000..4f99b3d --- /dev/null +++ b/keys/00 @@ -0,0 +1 @@ +ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCmJoiGO5YD3lbbIOJ99Al2xxm6QS9q+dTCTtlALjYI5f9ICGZJT8PEGlV9BBNCRQdgb3i2LBzQi90Tq1oG6/PcTV3Mto2TawLz5+2+ym29eIq1QIhVTLmZskK815FpawWqxY6+xpGU3vP1WjrFBbhGtl+CCaN+P2TWNkrR8FjG2144hdAlFfEEqfQC+TXbsyJCYoExuxGDJo8ae0JGbz9w1A1UbjnHwKnoxvirTFEbw9IHJIcTdUwuQKOrwydboCOqeaHt74+BnnCOZhpYqMDacrknHITN4GfFFzbs6FsE8NAwFk6yvkNXXzoe60iveNXtCIYuWjG517LQgHAC5BdaPgqzYNg+eqSul72e+jjRs+KDioNqvprw+TcBBO1lXZ2VQFyWyAdV2Foyaz3Wk5qYlOpX/9JLEp6H3cU0XCFR25FdXmjQ4oXN1QEe+2akV8MQ9cWhFhDcbY8Q1EiMWpBVC1xbt4FwE8VCTByZOZsQ0wPVe/vkjANOo+brS3tsR18= 00@xmuhpc diff --git a/keys/01 b/keys/01 new file mode 100644 index 0000000..8d9a5e0 --- /dev/null +++ b/keys/01 @@ -0,0 +1 @@ +ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCxcIWDQxVyIRqCGR4uWtrh4tLc025+q6du2GVsox8IzmBFkjNY8Au5GIMP5BKRstxFdg3f/wam8krckUN9rv5+OHB9U8HGz77Xs0FktqRVNMaDPdptePZQJ9A9eW3kkFDfQnORJtiVcEWfUBS3pi0QFOHylnG27YyC/Vjx9tjvtJWKsQEVTFJbFHPdi+G7lHTpqIGx+/a2JN9O6uVujXXYvjSVXsd+CWB9VMZMvYCIz2Ecb6RqR3brj4FhRRl8zyCj+J4ACYFdGWL98fTab2uPHbpVeKrefFFA43JOD/4zwBx/uw7MAQAq0GunTV3FpBfIAQHWgftf2fSlbz20oPjCwdYn9ZuGJOBUroryex7AKZmnSYM3biLHcctQfZtxqVPEU3W/62MUsI/kZb9RcF24JRksMoS2XWTiv2HFf5ijQGLXXOjqiTlGncwiKf65DwkDBsSxzgbXk5Uo86viq6UITFXPx/RytU+SUiN4Wb7wcBTjt/+tyQd1uqc7+3DCDXk= 01@xmuhpc diff --git a/keys/02 b/keys/02 new file mode 100644 index 0000000..bf6dce7 --- /dev/null +++ b/keys/02 @@ -0,0 +1 @@ +ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQDkT/P4MnzxBh8sRi0oQ88duNpY/ejFtptGqUQJVobj23vbu7ju6x/yuXqnHFOLi/IOZgNl5oBhRlJekRL+FWMIwpPBA6MnbVNkHXvwu5kLXVTt0O9dhJfDiPPbYcNjOhw4o8aZMc0oEyz8xZgkPoIehHQda+K5vRhFnYCRgn2X92VY/dW1QqPJKEfN47Tsp00w8wyKixEvuJe8OBEoKDpiZYzbXJKuoKhCdMp0uMHMCojYuYP9rGZO6bHl7Q6cYotGx1jH2pe30Ujtm3Xbm44H1mhXr1K/lhcHfojSge8POqii+eaXSCzqRlXaWyvrL9JLaaRD7GfWDaRWSKDfN8Ha4mnUvRtObRMSLOnr2QOTLJw9QPnlDDxCd1q7yluKraccYnTQQP5JuBwkRqjuJTatd9b18Z14HffmXZNR7asT1sJXK1rWKeLTrZwqxpkuwLAnbr60PVwfMHZeZ6FVPXGZ4wQb22lFHvaZZCEJf+9QDXpDn5L59FlaBYO2Xwojj3s= 02@xmuhpc diff --git a/keys/03 b/keys/03 new file mode 100644 index 0000000..d3d02a6 --- /dev/null +++ b/keys/03 @@ -0,0 +1 @@ +ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQDOF3LfnQiI8wpsXGn87bt7rbUZcgsdaOSOswk4Vf4dBautEdQZc0q+UDB2TlR2K8L7SPyywpl5z67euN5QRJLEwg8flTybiJp3EKDctYEM22sa36ONcSIJ/iHSdCkwtPXkBYreh9e+MAHfTroIKK5zM/P1QIN3NrknIXpWjLDF73ejrxE+EXRK6jbuWfo+5dnLnDoUFt1e+pYLZos5KRRB94Qt5I79D/cAg3hG+Zl2FCCOpn1hIdLo/kWJTKUPe61oUaIxriV6nCXp/pU1BHlM43hGowiHa4bVZIs8Eo4r7OI9thhSuS2BKSifibBKIicZtntSlS/I3xa5am28YLmrOiEXRsjPom7trO8qIhPfYOc/yFDg1gcpLxyNroCPooPBzPxUqrTT96Q4fDDTaqfyuVxQFxbYoFAqQs8/lw6WcGJ4fGC5JPsPiwoSdQy/B7gCfQcFjPXp1NH8Sx+xMLCmxRqdKSyeiEwoyB0tZ6ngaI73HFhCPX1/rLx3xv0zd/8= 03@xmuhpc diff --git a/keys/chn b/keys/chn new file mode 100644 index 0000000..348493e --- /dev/null +++ b/keys/chn @@ -0,0 +1 @@ +ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDXlhoouWG+arWJz02vBP/lxpG2tUjx8jhGBnDeNyMu0OtGcnHMAWcb3YDP0A2XJIVFBCCZMM2REwnSNbHRSCl1mTdRbelfjA+7Jqn1wnrDXkAOG3S8WYXryPGpvavu6lgW7p+dIhGiTLWwRbFH+epFTn1hZ3A1UofVIWTOPdoOnx6k7DpQtIVMWiIXLg0jIkOZiTMr3jKfzLMBAqQ1xbCV2tVwbEY02yxxyxIznbpSPReyn1RDLWyqqLRd/oqGPzzhEXNGNAZWnSoItkYq9Bxh2AvMBihiTir3FEVPDgDLtS5LUpM93PV1yTr6JyCPAod9UAxpfBYzHKse0KCQFoZH chn@chn-PC diff --git a/keys/gb b/keys/gb new file mode 100644 index 0000000..4385849 --- /dev/null +++ b/keys/gb @@ -0,0 +1 @@ +ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQDWAfyfDFctbzJTiuK9IPw3yFLqt7vqd/T0/HoZfH/bzLZ8GVeod2oz6kjm3ns0IG94HO5vGMEmQfbK1ZKT2TqA7ve+3wG9seiwfh8xh7Xhl2AnaF0pjHEXnw+w8mTzxCv9qRhsgfHuuBVhH6PguHvk66GKjvNaxTJhlKAyNogOI3jLnw7ODFScldHbJlMYl1pBHV/G/Zeuq0qnA/pkeiFdvlsZUVGD0cCfuoHm8FCfEzv6pfkhVJUH0v5rof8GiT9eg7ntG49Gei1lkH5NosbY8f6fEKNSoOc0dm5g2FaI3D7LJixwQ6rMiJwmPb6A4oHmcJQKokU8uhROQorYLgV7RtrnHu2cHMRW6SiAUvpmvaPPcxn8CbfuSOGDhYRKxNJNtWRK08Urtq9tYD+Fpze4QoZXxN35uvsi3lMA55PK0AsTm/aVGslzHUUzgWtDxcI2pLAm9rFpCRPCY+UC1Xp5vjZoqZXwhJ81qZ7VXWTM2voxCrKAlu+Zg2FaQD5szOU= gb@xmupc1 diff --git a/keys/xll b/keys/xll new file mode 100644 index 0000000..f9f55de --- /dev/null +++ b/keys/xll @@ -0,0 +1 @@ +ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQDJ/jzUQ6QuAjnAryvpWk7TReS6pnHxhEXY9RonojKkurhfYSQO/IlxDMDq23TFXcgu8iZG4cS6MADgx/KNZD/MjuN9YNCIEGvMwzWvB0oM25BC6Vf3iKDmhH06rZKH6/g5GN+HWoCN4yE/+MhIpegFO3+YMpveXwEESlyoIjPvcW+RwmlNJevrHd83ETYDQ4AybWyJo6en5tz2ngr22HaK4MtxgrqnIN/KorY+nrzTNa7VBC7BaZc1tA5FLwUeCXtuzp2ibfrxoGUAiDig4FW09ijCk3Y77y7aNVI2nw5y28nCV5rgVMh5fejtNVqIqku7p+8qgjxvY6veATG0lYgZgw2ldnDGDNbEGxcCnKKmCgZMxok8zTRsniZ91KuHkcl2L7xUo7kdQYzBRwZyQ53eW+yPoqUya4yn272rscBEUMyZzmegfr1SXMqw/8zn+MZdr1KXEvrbfjX+2QL52GY3bfYUf3KFje+Sp88k688bRH0vrxj9BCOS7ovbyfe9BEU= xll@xmupc1 diff --git a/keys/yjq b/keys/yjq new file mode 100644 index 0000000..2e62e5d --- /dev/null +++ b/keys/yjq @@ -0,0 +1 @@ +ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCtnVhZQsJfbs2w9hFZkx4qDhIs++7no+6r5TifP3Dq7epJYd2QYx4dI66XxTNhKxZjN6a4Xn5nFlYLtQJXOvzBLC8IBf1W5GCH0k/jqzzskS0/Ix/70HzcBwJk8ihWDkyON5Ki1BRCx34RNxth1BIxWyc5QT+lou+D92x8iAu/uOvmcAL3Ua0OlZwxw03hLp/PpS4ZnUqFjc2JVtarY7eQu/i3RwOZUaK6nT2EL8RObzk4xnieqsU5PWwA3voVjetqZaDQ+P7dimQXz/FaucroKxCNyTiy1oG4fdQpm2UDrH6ZfPvdQLYrtet6FQabXOxhV7MuR3jYtxZjs1kDVZIseIZ6IwjetaUoMxvIouRfYjOSIEo9Ek9o0+Yhku4r0uWmPDrymWugU1raMmlRxSUwdlzW+C7mQwtGbs/MG4MN4GWkM6id5DKlY2vYKUfrTzmhY1swCtzKq20fjvyX8qhJdcytgVlOrBZnPje6Qd55sI0RjdgJrBsxT2SYquez7U8= yjq@xmupc1 diff --git a/keys/zem b/keys/zem new file mode 100644 index 0000000..65a5151 --- /dev/null +++ b/keys/zem @@ -0,0 +1 @@ +ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQDn1pfGen7kjPTHsbb8AgrUJWOeFPHK5S4M97Lcj3tvdcjZi2SXN6PwHQfh8/xGhZbTLPz/40S9O9/Dn30xkUTfnONirKt790jp7VEbOtPnjQPOd/KRNWlS3VV0BELuq5p633Mi13rP6JZtdKmU2uSkvvaUBfCppy3JaWv/B7HLJ48f8IzkdiT1px3dN1eQ4SFoHOiVG0ci5TGG6wfMdoAAnM9R1aXI4gDxnYjLYujpaNZ4hBOta/6ZK/PV0JufoXdIAZjubgk1Hv04XHXLR2Z0UhRM6x7UrZIOdM/LlnKmcVk408ZKEj/9m1xRyDsNoZ24CF++cmnwfBHrp9I5nvDI7xOTdZlOhzkiiPM3f4i6s2Qjdv4vpZ6AeE3Qt1LVQyAr67b4UMjHuYqSi2KgyCO6My2Ov2eRoS74EKcb8ejJv3O+XInmYUgDgTgDFT3CgQgK2DG45HiV6nOkaE/6iKx2JSOiYZTFc7TRcePfXF9JQD7dXFde6qm3EbIVyJIpCJ8= zem@xmupc1 diff --git a/src/common.cpp b/src/common.cpp new file mode 100644 index 0000000..3e5060a --- /dev/null +++ b/src/common.cpp @@ -0,0 +1,35 @@ +# include +# include +# include +# include + +namespace hpcstat +{ + std::optional exec + (std::filesystem::path program, std::vector args, std::optional stdin) + { + namespace bp = boost::process; + bp::ipstream output; + bp::opstream input; + std::unique_ptr process; + if (stdin) + { + process = std::make_unique + (program.string(), bp::args(args), bp::std_out > output, bp::std_err > stderr, bp::std_in < input); + input << *stdin; + input.close(); + } + else process = std::make_unique + (program.string(), bp::args(args), bp::std_out > output, bp::std_err > stderr, bp::std_in < bp::null); + process->wait(); + if (process->exit_code() != 0) return std::nullopt; + std::stringstream ss; + ss << output.rdbuf(); + return ss.str(); + } + long now() + { + return std::chrono::duration_cast + (std::chrono::system_clock::now().time_since_epoch()).count(); + } +} diff --git a/src/env.cpp b/src/env.cpp new file mode 100644 index 0000000..5eb9024 --- /dev/null +++ b/src/env.cpp @@ -0,0 +1,19 @@ +# include +# include +# include +# include + +namespace hpcstat::env +{ + bool interactive() { return isatty(fileno(stdin)); } + std::optional env(std::string name, bool required) + { + if (auto value = std::getenv(name.c_str()); !value) + { + if (required) std::cerr << fmt::format("Failed to get environment variable {}\n", name); + return std::nullopt; + } + else return value; + } + // XDG_SESSION_ID HPCSTAT_SUBACCOUNT SSH_CONNECTION +} diff --git a/src/keys.cpp b/src/keys.cpp new file mode 100644 index 0000000..5e3b21d --- /dev/null +++ b/src/keys.cpp @@ -0,0 +1,17 @@ +# include + +namespace hpcstat +{ + std::map Keys + { + { "LNoYfq/SM7l8sFAy325WpC+li+kZl3jwST7TmP72Tz8", { "chn", "Haonan Chen" } }, + { "VJT5wgkb2RcIeVNTA+/NKxokctbYnJ/KgH6IxrKqIGE", { "gb", "Bin Gong" } }, + { "umC3/RB1vS8TQBHsY3IzhOiyqVrOSw2fB3rIpDQSmf4", { "xll", "Leilei Xiang" } }, + { "fdq5k13N2DAzIK/2a1Mm4/ZVsDUgT623TSOXsVswxT8", { "yjq", "Junqi Yao" } }, + { "8USxEYi8ePPpLhk5FYBo2udT7/NFmEe8c2+oQajGXzA", { "zem", "Enming Zhang" } }, + { "7bmG24muNsaAZkCy7mQ9Nf2HuNafmvUO+Hf1bId9zts", { "00", "Yaping Wu" } }, + { "dtx0QxdgFrXn2SYxtIRz43jIAH6rLgJidSdTvuTuews", { "01", "Jing Li" } }, + { "8crUO9u4JiVqw3COyjXfzZe87s6XZFhvi0LaY0Mv6bg", { "02", "Huahan Zhan" } }, + { "QkmIYw7rmDEAP+LDWxm6L2/XLnAqTwRUB7B0pxYlOUs", { "03", "Na Gao" } } + }; +} diff --git a/src/lfs.cpp b/src/lfs.cpp new file mode 100644 index 0000000..d8c7d12 --- /dev/null +++ b/src/lfs.cpp @@ -0,0 +1,100 @@ +# include +# include +# include +# include +# include +# include +# include +# include +# include + +namespace hpcstat::lfs +{ + std::optional> bsub(std::vector args) + { + if (auto bsub = env::env("HPCSTAT_BSUB", true); !bsub) + return std::nullopt; + else + { + std::set valid_args = { "J" "q" "n" "R" "o" }; + for (auto it = args.begin(); it != args.end(); it++) + { + if (it->length() > 0 && (*it)[0] == '-') + { + if (!valid_args.contains(it->substr(1))) + { + std::cerr << fmt::format("Unknown bsub argument: {}\n", *it) + << "bsub might support this argument, but hpcstat currently does not support it.\n" + "If you are sure this argument is supported by bsub,\n" + "please submit issue on [github](https://github.com/CHN-beta/hpcstat) or contact chn@chn.moe.\n"; + return std::nullopt; + } + if (it + 1 != args.end() && ((it + 1)->length() == 0 || (*(it + 1))[0] != '-')) it++; + } + else break; + } + if (auto result = exec(*bsub, args); result) return std::nullopt; + else + { + // Job <462270> is submitted to queue . + std::regex re(R"r(Job <(\d+)> is submitted to queue <(\w+)>.)r"); + std::smatch match; + if (std::regex_search(*result, match, re)) + return std::make_pair(std::stoi(match[1]), match[2]); + else + { + std::cerr << fmt::format("Failed to parse job id from output: {}\n", *result); + return std::nullopt; + } + } + } + } + std::optional>> bjobs_list() + { + if + ( + auto result = exec + ( + boost::process::search_path("bjobs").string(), + { "-a", "-o", "jobid submit_time stat cpu_used", "-json" } + ); + !result + ) + return std::nullopt; + else + { + nlohmann::json j; + try { j = nlohmann::json::parse(*result); } + catch (nlohmann::json::parse_error& e) + { + std::cerr << fmt::format("Failed to parse bjobs output: {}\n", e.what()); + return std::nullopt; + } + std::map> jobs; + for (auto& job : j["RECORDS"]) + { + std::string status = job["STAT"]; + if (!std::set{ "DONE", "EXIT" }.contains(status)) continue; + std::string submit_time = job["SUBMIT_TIME"]; + std::string cpu_used_str = job["CPU_USED"]; + double cpu_used = std::stof(cpu_used_str.substr(0, cpu_used_str.find(' '))); + jobs[std::stoi(job["JOBID"].get())] = { submit_time, status, cpu_used }; + } + return jobs; + } + } + std::optional bjobs_detail(unsigned jobid) + { + if + ( + auto result = exec + ( + boost::process::search_path("bjobs").string(), + { "-l", std::to_string(jobid) } + ); + !result + ) + return std::nullopt; + else return *result; + } +} diff --git a/src/main.cpp b/src/main.cpp index fabf6a7..2ea3af6 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -1,213 +1,112 @@ -# include -# include -# include -# include -# include -# include -# include -# include -# include -# include -# include -# include +# include +# include +# include +# include +# include +# include # include -# include -# include - -using namespace std::literals; - -// ssh fingerprint -> username -const std::map Username -{ - { "LNoYfq/SM7l8sFAy325WpC+li+kZl3jwST7TmP72Tz8", "Haonan Chen" }, - { "VJT5wgkb2RcIeVNTA+/NKxokctbYnJ/KgH6IxrKqIGE", "Bin Gong" }, - { "umC3/RB1vS8TQBHsY3IzhOiyqVrOSw2fB3rIpDQSmf4", "Leilei Xiang" }, - { "fdq5k13N2DAzIK/2a1Mm4/ZVsDUgT623TSOXsVswxT8", "Junqi Yao" }, - { "8USxEYi8ePPpLhk5FYBo2udT7/NFmEe8c2+oQajGXzA", "Enming Zhang" }, - { "7bmG24muNsaAZkCy7mQ9Nf2HuNafmvUO+Hf1bId9zts", "Yaping Wu" }, - { "dtx0QxdgFrXn2SYxtIRz43jIAH6rLgJidSdTvuTuews", "Jing Li" }, - { "8crUO9u4JiVqw3COyjXfzZe87s6XZFhvi0LaY0Mv6bg", "Huahan Zhan" }, - { "QkmIYw7rmDEAP+LDWxm6L2/XLnAqTwRUB7B0pxYlOUs", "Na Gao" } -}; - -// program path, set at start of main, e.g. /gpfs01/.../bin/hpcstat -std::filesystem::path Program; - -// run a program, wait until it exit, return its stdout if it return 0, otherwise nullopt -std::optional exec(boost::filesystem::path program, std::vector args) -{ - namespace bp = boost::process; - bp::ipstream output; - auto process = bp::child - (program, bp::args(args), bp::std_out > output, bp::std_err > stderr, bp::std_in < bp::null); - process.wait(); - if (process.exit_code() != 0) return std::nullopt; - std::stringstream ss; - ss << output.rdbuf(); - return ss.str(); -} - -// detect ssh fingerprint using ssh-add -// always assume sha256 fingerprint -std::optional> fingerprints() -{ - auto output = - exec(Program.replace_filename("ssh-add"), { "-l" }); - if (!output) { std::cerr << "Failed to get ssh fingerprints\n"; return std::nullopt; } - auto fingerprint = output->substr(0, 47); - // search for all strings that match the fingerprint pattern: sha256:... - std::regex pattern(R"r(\b(?:sha|SHA)256:([0-9A-Za-z+/=]{43})\b)r"); - std::smatch match; - std::vector fingerprints; - for (auto i = std::sregex_iterator(output->begin(), output->end(), pattern); i != std::sregex_iterator(); i++) - fingerprints.push_back(i->str(1)); - return fingerprints; -} - -// get an authenticated fingerprint and username -std::optional> authenticated() -{ - auto fps = fingerprints(); - if (!fps) return std::nullopt; - for (auto& fp : *fps) - if (Username.contains(fp)) return std::make_pair(fp, Username.at(fp)); - std::cerr << fmt::format("No valid fingerprint found, available fingerprints: {}\n", *fps); - return std::nullopt; -} - -// initialize the database -struct LoginData -{ - unsigned Id = 0; long Time = 0; - std::string Key, SessionId; - std::optional Subaccount; - bool Interactive; -}; -using LoginTable = zxorm::Table -< - "login", LoginData, - zxorm::Column<"id", &LoginData::Id, zxorm::PrimaryKey<>>, - zxorm::Column<"time", &LoginData::Time>, - zxorm::Column<"key", &LoginData::Key>, - zxorm::Column<"session_id", &LoginData::SessionId>, - zxorm::Column<"sub_account", &LoginData::Subaccount>, - zxorm::Column<"interactive", &LoginData::Interactive> ->; -struct LogoutData { unsigned Id = 0; long Time = 0; std::string SessionId; }; -using LogoutTable = zxorm::Table -< - "logout", LogoutData, - zxorm::Column<"id", &LogoutData::Id, zxorm::PrimaryKey<>>, - zxorm::Column<"time", &LogoutData::Time>, zxorm::Column<"sessionid", &LogoutData::SessionId> ->; -struct SubmitJobData - { unsigned Id = 0; long Time = 0; int JobId; std::string Key, SessionId, Subaccount, SubmitDir, JobCommand; }; -using SubmitJobTable = zxorm::Table -< - "submitjob", SubmitJobData, - zxorm::Column<"id", &SubmitJobData::Id, zxorm::PrimaryKey<>>, - zxorm::Column<"time", &SubmitJobData::Time>, - zxorm::Column<"job_id", &SubmitJobData::JobId>, - zxorm::Column<"key", &SubmitJobData::Key>, - zxorm::Column<"session_id", &SubmitJobData::SessionId>, - zxorm::Column<"sub_account", &SubmitJobData::Subaccount>, - zxorm::Column<"submit_dir", &SubmitJobData::SubmitDir>, - zxorm::Column<"job_command", &SubmitJobData::JobCommand> ->; -struct FinishJobData { unsigned Id = 0; long Time = 0; int JobId; std::string JobResult; double CpuTime; }; -using FinishJobTable = zxorm::Table -< - "finishjob", FinishJobData, - zxorm::Column<"id", &FinishJobData::Id, zxorm::PrimaryKey<>>, - zxorm::Column<"time", &FinishJobData::Time>, - zxorm::Column<"job_id", &FinishJobData::JobId>, - zxorm::Column<"job_result", &FinishJobData::JobResult>, - zxorm::Column<"cpu_time", &FinishJobData::CpuTime> ->; -struct QueryJobData { unsigned Id = 0; int JobId; }; -using QueryJobTable = zxorm::Table -< - "queryjob", QueryJobData, - zxorm::Column<"id", &QueryJobData::Id, zxorm::PrimaryKey<>>, zxorm::Column<"job_id", &QueryJobData::JobId> ->; -void initdb() -{ - auto dbfile = Program.replace_filename("hpcstat.db").string(); - zxorm::Connection - conn(dbfile.c_str()); - conn.create_tables(); -} -void writedb(auto value) -{ - auto dbfile = Program.replace_filename("hpcstat.db").string(); - zxorm::Connection - conn(dbfile.c_str()); - value.Time = std::chrono::duration_cast - (std::chrono::system_clock::now().time_since_epoch()).count(); - conn.insert_record(value); -} - -bool interactive() { return isatty(fileno(stdin)); } - -// get value of XDG_SESSION_ID -std::optional session_id() -{ - if (auto value = std::getenv("XDG_SESSION_ID"); !value) - { std::cerr << "Failed to get session id\n"; return std::nullopt; } - else return value; -} - -// get value of HPCSTAT_SUBACCOUNT -std::optional subaccount() - { if (auto value = std::getenv("HPCSTAT_SUBACCOUNT"); value) return value; else return std::nullopt; } +# include int main(int argc, const char** argv) { + using namespace hpcstat; std::vector args(argv, argv + argc); - Program = boost::dll::program_location().string(); if (args.size() == 1) { std::cout << "Usage: hpcstat initdb|login|logout|submitjob|finishjob\n"; return 1; } else if (args[1] == "initdb") - initdb(); + { + if (!sql::initdb()) { std::cerr << "Failed to initialize database\n"; return 1; } + } else if (args[1] == "login") { - std::cout << "Checking your ssh certification..." << std::flush; - if (auto key = authenticated(); !key) return 1; - else if (auto session = session_id(); !session) return 1; + std::cout << "Communicating with the agent..." << std::flush; + if (auto fp = ssh::fingerprint(); !fp) return 1; + else if (auto session = env::env("XDG_SESSION_ID", true); !session) + return 1; else { - writedb(LoginData - {.Key = key->first, .SessionId = *session, .Subaccount = subaccount(), .Interactive = interactive()}); - std::cout << fmt::format("\33[2K\rLogged in as {}.\n", key->second); + sql::LoginData data + { + .Time = now(), .Key = *fp, .SessionId = *session, .Subaccount = env::env("HPCSTAT_SUBACCOUNT"), + .Ip = env::env("SSH_CONNECTION"), .Interactive = env::interactive() + }; + auto signature = ssh::sign(sql::serialize(data), *fp); + if (!signature) return 1; + data.Signature = *signature; + sql::writedb(data); + std::cout << fmt::format("\33[2K\rLogged in as {}.\n", *fp); } } else if (args[1] == "logout") { - if (auto session = session_id(); !session) return 1; - else writedb(LogoutData{.SessionId = *session}); + if (auto session_id = env::env("XDG_SESSION_ID", true); !session_id) + return 1; + else sql::writedb(sql::LogoutData{ .Time = now(), .SessionId = *session_id }); } else if (args[1] == "submitjob") { - if (args.size() < 4) { std::cerr << "Usage: hpcstat submitjob \n"; return 1; } - if (auto key = authenticated(); !key) return 1; - else if (auto session = session_id(); !session) return 1; - else writedb(SubmitJobData + if (args.size() < 3) { std::cerr << "Usage: hpcstat submitjob \n"; return 1; } + if (auto fp = ssh::fingerprint(); !fp) return 1; + else if (auto session = env::env("XDG_SESSION_ID", true); !session) + return 1; + else if + (auto bsub = lfs::bsub(args | ranges::views::drop(2) | ranges::to>); !bsub) + return 1; + else { - .JobId = std::stoi(args[2]), .Key = key->first, .SessionId = *session, - .SubmitDir = std::filesystem::current_path().string(), - .JobCommand = [&] - { std::stringstream ss; for (int i = 3; i < args.size(); i++) ss << args[i] << " "; return ss.str(); }() - }); + sql::SubmitJobData data + { + .Time = now(), .JobId = bsub->first, .Key = *fp, .SessionId = *session, + .SubmitDir = std::filesystem::current_path().string(), + .JobCommand = args | ranges::views::drop(2) | ranges::views::join(' ') | ranges::to(), + .Subaccount = env::env("HPCSTAT_SUBACCOUNT"), .Ip = env::env("SSH_CONNECTION") + }; + auto signature = ssh::sign(sql::serialize(data), *fp); + if (!signature) return 1; + data.Signature = *signature; + sql::writedb(data); + std::cout << fmt::format + ("Job {} was submitted to {} by {}.\n", bsub->first, bsub->second, Keys[*fp].Username); + } } else if (args[1] == "finishjob") { - - } - else - { - std::cerr << "Unknown command\n"; - return 1; + if (auto fp = ssh::fingerprint(); !fp) return 1; + else if (auto session = env::env("XDG_SESSION_ID", true); !session) + return 1; + else if (auto all_jobs = lfs::bjobs_list(); !all_jobs) return 1; + else if + ( + auto not_recorded = sql::finishjob_remove_existed + ( + *all_jobs + | ranges::views::transform([](auto& it) { return std::pair{ it.first, std::get<0>(it.second) }; }) + | ranges::to> + ); + !not_recorded + ) + return 1; + else for (auto jobid : *not_recorded) + { + if (auto detail = lfs::bjobs_detail(jobid); !detail) return 1; + else + { + sql::FinishJobData data + { + .Time = now(), .JobId = jobid, .JobResult = std::get<1>(all_jobs->at(jobid)), + .SubmitTime = std::get<1>(all_jobs->at(jobid)), .CpuTime = std::get<2>(all_jobs->at(jobid)), + }; + if + ( + auto signature = ssh::sign(sql::serialize(data), *fp); + !signature + ) + return 1; + else { data.Signature = *signature; sql::writedb(data); } + } + } } + else { std::cerr << "Unknown command.\n"; return 1; } return 0; } diff --git a/src/sql.cpp b/src/sql.cpp new file mode 100644 index 0000000..956e5a0 --- /dev/null +++ b/src/sql.cpp @@ -0,0 +1,55 @@ +# include +# include +# include +# include +# include +# include + +namespace hpcstat::sql +{ + std::string serialize(auto data) + { + auto [serialized_data_byte, out] = zpp::bits::data_out(); + out(data).or_throw(); + static_assert(sizeof(char) == sizeof(std::byte)); + return { reinterpret_cast(serialized_data_byte.data()), serialized_data_byte.size() }; + } + template std::string serialize(LoginData); + template std::string serialize(SubmitJobData); + template std::string serialize(FinishJobData); + std::optional> connect() + { + if (auto datadir = env::env("HPCSTAT_DATADIR", true); !datadir) + return std::nullopt; + else + { + auto dbfile = std::filesystem::path(*datadir) / "hpcstat.db"; + return std::make_optional> + (dbfile.c_str()); + } + } + bool initdb() + { if (auto conn = connect(); !conn) return false; else { conn->create_tables(); return true; } } + bool writedb(auto value) + { if (auto conn = connect(); !conn) return false; else { conn->insert_record(value); return true; } } + template bool writedb(LoginData); + template bool writedb(LogoutData); + template bool writedb(SubmitJobData); + template bool writedb(FinishJobData); + std::optional> finishjob_remove_existed(std::map jobid_submit_time) + { + if (auto conn = connect(); !conn) return std::nullopt; + else + { + auto all_job = jobid_submit_time | ranges::views::keys | ranges::to>; + auto not_logged_job = all_job | ranges::to>; + for (auto it : conn->select_query() + .order_by>(zxorm::order_t::DESC) + .where_many(FinishJobTable::field_t<"id">().in(all_job)) + .exec()) + if (jobid_submit_time[it.JobId] == it.SubmitTime) + not_logged_job.erase(it.JobId); + return not_logged_job; + } + } +} diff --git a/src/ssh.cpp b/src/ssh.cpp new file mode 100644 index 0000000..0ad6bf2 --- /dev/null +++ b/src/ssh.cpp @@ -0,0 +1,84 @@ +# include +# include +# include +# include +# include +# include +# include +# include +# include +# include +# include + +namespace hpcstat::ssh +{ + std::optional fingerprint() + { + if (auto datadir = env::env("HPCSTAT_DATADIR", true); !datadir) + return std::nullopt; + else if + ( + auto output = + exec(std::filesystem::path(*datadir) / "ssh-add", { "-l" }); + !output + ) + { std::cerr << "Failed to get ssh fingerprints\n"; return std::nullopt; } + else + { + std::regex pattern(R"r(\b(?:sha|SHA)256:([0-9A-Za-z+/=]{43})\b)r"); + std::smatch match; + for + ( + auto i = std::sregex_iterator(output->begin(), output->end(), pattern); + i != std::sregex_iterator(); i++ + ) + if (Keys.contains(i->str(1))) return i->str(1); + std::cerr << fmt::format("No valid fingerprint found in:\n{}\n", *output); + return std::nullopt; + } + } + std::optional sign(std::string message, std::string fingerprint) + { + if (auto datadir = env::env("HPCSTAT_DATADIR", true); !datadir) + return std::nullopt; + else if + ( + auto output = exec + ( + std::filesystem::path(*datadir) / "ssh-keygen", + { + "-Y", "sign", + "-f", fmt::format("{}/keys/{}", *datadir, Keys[fingerprint].PubkeyFilename), + "-n", "hpcstat@chn.moe", "-" + }, + message + ); + !output + ) + { std::cerr << fmt::format("Failed to sign message: {}\n", message); return std::nullopt; } + else return *output; + } + bool verify(std::string message, std::string signature, std::string fingerprint) + { + if (auto datadir = env::env("HPCSTAT_DATADIR", true); !datadir) + return false; + else + { + namespace bf = boost::filesystem; + auto tempdir = bf::temp_directory_path() / bf::unique_path(); + bf::create_directories(tempdir); + auto signaturefile = tempdir / "signature"; + std::ofstream(signaturefile) << signature; + return exec + ( + std::filesystem::path(*datadir) / "ssh-keygen", + { + "-Y", "verify", + "-f", fmt::format("{}/keys/{}", *datadir, Keys[fingerprint].PubkeyFilename), + "-n", "hpcstat@chn.moe", "-s", signaturefile.string() + }, + message + ).has_value(); + } + } +}