gbf_suite/
aws_upload.rs

1use aws_config::BehaviorVersion;
2use aws_sdk_dynamodb::types::AttributeValue;
3use serde_dynamo::aws_sdk_dynamodb_1::to_attribute_value;
4
5use crate::{
6    consts::{self, GBF_AWS_BUCKET},
7    gbf_result::{
8        GbfFunctionDao, GbfFunctionErrorDao, GbfGraphvizStructureAnalaysisDao, GbfModuleDao,
9        GbfVersionDao,
10    },
11    utils::hash_string,
12};
13
14pub struct AwsUpload {
15    pub s3_client: aws_sdk_s3::Client,
16    pub dynamo_client: aws_sdk_dynamodb::Client,
17}
18
19impl AwsUpload {
20    pub async fn new() -> Self {
21        // Load AWS credentials
22        let sdk_config = aws_config::load_defaults(BehaviorVersion::latest()).await;
23        let s3_client = aws_sdk_s3::Client::new(&sdk_config);
24        let dynamo_client = aws_sdk_dynamodb::Client::new(&sdk_config);
25        Self {
26            s3_client,
27            dynamo_client,
28        }
29    }
30
31    pub async fn upload_graphviz_dot(
32        &self,
33        graphviz_dot: String,
34    ) -> Result<String, aws_sdk_s3::Error> {
35        // Upload the file to S3
36        let s3_key = format!("graphviz/{}.dot", hash_string(&graphviz_dot));
37
38        self.s3_client
39            .put_object()
40            .bucket(consts::GBF_AWS_BUCKET)
41            .content_type("text/plain")
42            .key(s3_key.clone())
43            .body(graphviz_dot.into_bytes().into())
44            .send()
45            .await?;
46        Ok(s3_key)
47    }
48
49    pub async fn upload_gbf_version(
50        &self,
51        gbf_version: GbfVersionDao,
52    ) -> Result<(), aws_sdk_dynamodb::Error> {
53        self.dynamo_client
54            .put_item()
55            .table_name(consts::GBF_AWS_DYNAMO_VERSION_TABLE)
56            .item(
57                gbf_version.pk_key(),
58                AttributeValue::S(gbf_version.pk_val()),
59            )
60            .item("gbf_version", AttributeValue::S(gbf_version.gbf_version))
61            .item(
62                "total_time",
63                AttributeValue::N(gbf_version.total_time.as_millis().to_string()),
64            )
65            .item(
66                "suite_timestamp",
67                AttributeValue::N(gbf_version.suite_timestamp.to_string()),
68            )
69            .send()
70            .await?;
71        Ok(())
72    }
73
74    pub async fn upload_gbf_module(
75        &self,
76        gbf_module: GbfModuleDao,
77    ) -> Result<(), aws_sdk_dynamodb::Error> {
78        self.dynamo_client
79            .put_item()
80            .table_name(consts::GBF_AWS_DYNAMO_MODULE_TABLE)
81            .item(gbf_module.pk_key(), AttributeValue::S(gbf_module.pk_val()))
82            .item("gbf_version", AttributeValue::S(gbf_module.gbf_version))
83            .item("module_id", AttributeValue::S(gbf_module.module_id))
84            .item("file_name", AttributeValue::S(gbf_module.file_name))
85            .item(
86                "module_load_time",
87                AttributeValue::N(gbf_module.module_load_time.as_millis().to_string()),
88            )
89            .item(
90                "decompile_success",
91                AttributeValue::Bool(gbf_module.decompile_success),
92            )
93            .send()
94            .await?;
95        Ok(())
96    }
97
98    pub async fn upload_gbf_function(
99        &self,
100        gbf_function: GbfFunctionDao,
101    ) -> Result<(), aws_sdk_dynamodb::Error> {
102        self.dynamo_client
103            .put_item()
104            .table_name(consts::GBF_AWS_DYNAMO_FUNCTION_TABLE)
105            .item(
106                gbf_function.pk_key(),
107                AttributeValue::S(gbf_function.pk_val()),
108            )
109            .item(
110                "gbf_version",
111                AttributeValue::S(gbf_function.clone().gbf_version),
112            )
113            .item(
114                "module_id",
115                AttributeValue::S(gbf_function.clone().module_id),
116            )
117            .item(
118                "function_address",
119                AttributeValue::N(gbf_function.clone().function_address.to_string()),
120            )
121            .item(
122                "function_name",
123                AttributeValue::S(
124                    gbf_function
125                        .clone()
126                        .function_name
127                        .unwrap_or("entry".to_string()),
128                ),
129            )
130            .item(
131                "decompile_success",
132                AttributeValue::Bool(gbf_function.clone().decompile_success),
133            )
134            .item(
135                "decompile_result",
136                AttributeValue::S(
137                    gbf_function
138                        .clone()
139                        .decompile_result
140                        .unwrap_or("".to_string()),
141                ),
142            )
143            .item(
144                "total_time",
145                AttributeValue::N(gbf_function.clone().total_time.as_millis().to_string()),
146            )
147            .item(
148                "dot_url",
149                AttributeValue::S(gbf_function.dot_url(GBF_AWS_BUCKET)),
150            )
151            .send()
152            .await?;
153        Ok(())
154    }
155
156    pub async fn upload_gbf_function_error(
157        &self,
158        gbf_function_error: GbfFunctionErrorDao,
159    ) -> Result<(), Box<dyn std::error::Error>> {
160        let backtrace = serde_json::to_value(&gbf_function_error.backtrace)?;
161        let attr = to_attribute_value(backtrace)?;
162
163        let context = serde_json::to_value(&gbf_function_error.context)?;
164        let attr_ctx = to_attribute_value(context)?;
165
166        self.dynamo_client
167            .put_item()
168            .table_name(consts::GBF_AWS_DYNAMO_FUNCTION_ERROR_TABLE)
169            .item(
170                gbf_function_error.pk_key(),
171                AttributeValue::S(gbf_function_error.pk_val()),
172            )
173            .item(
174                "gbf_version",
175                AttributeValue::S(gbf_function_error.gbf_version),
176            )
177            .item("module_id", AttributeValue::S(gbf_function_error.module_id))
178            .item(
179                "function_address",
180                AttributeValue::N(gbf_function_error.function_address.to_string()),
181            )
182            .item(
183                "error_type",
184                AttributeValue::S(gbf_function_error.error_type),
185            )
186            .item("message", AttributeValue::S(gbf_function_error.message))
187            .item("backtrace", attr)
188            .item("context", attr_ctx)
189            .send()
190            .await?;
191        Ok(())
192    }
193
194    pub async fn upload_gbf_graphviz_dao(
195        &self,
196        graphviz: GbfGraphvizStructureAnalaysisDao,
197    ) -> Result<(), aws_sdk_dynamodb::Error> {
198        self.dynamo_client
199            .put_item()
200            .table_name(consts::GBF_AWS_DYNAMO_GRAPHVIZ_TABLE)
201            .item(graphviz.pk_key(), AttributeValue::S(graphviz.pk_val()))
202            .item(
203                "gbf_version",
204                AttributeValue::S(graphviz.clone().gbf_version),
205            )
206            .item("module_id", AttributeValue::S(graphviz.clone().module_id))
207            .item(
208                "function_address",
209                AttributeValue::N(graphviz.function_address.to_string()),
210            )
211            .item(
212                "structure_analysis_step",
213                AttributeValue::N(graphviz.structure_analysis_step.to_string()),
214            )
215            .item(
216                "dot_url",
217                AttributeValue::S(graphviz.dot_url(GBF_AWS_BUCKET)),
218            )
219            .send()
220            .await?;
221        Ok(())
222    }
223}