patoka/worker/
plugin.rs

1use serde_derive::{Deserialize, Serialize};
2use serde_json;
3use serde_json::json;
4use std::collections::{HashMap};
5use std::fmt;
6
7use crate::core::env::{self, *};
8use crate::core::proxy;
9use crate::core::user_agent;
10use crate::worker::worker_message::{WorkerMessage, Dest, WorkerMessagePayload};
11
12#[derive(Clone, PartialEq, Copy, Serialize, Deserialize)]
13#[serde(rename_all = "snake_case")]
14pub enum WorkerPlugin {
15    Basic,
16    HeadlessBrowser,
17    None,
18}
19
20impl WorkerPlugin {
21    pub fn as_str(ws: WorkerPlugin) -> &'static str {
22        match ws {
23            WorkerPlugin::Basic => "basic",
24            WorkerPlugin::HeadlessBrowser => "headless_browser",
25            WorkerPlugin::None => "none",
26        }
27    }
28
29    pub fn from_str(s: &str) -> WorkerPlugin {
30        match s {
31            "basic" => WorkerPlugin::Basic,
32            "headless_browser" => WorkerPlugin::HeadlessBrowser,
33            _ => WorkerPlugin::None,
34        }
35    }
36}
37
38impl fmt::Debug for WorkerPlugin {
39    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
40        write!(f, "{:?}", WorkerPlugin::as_str(*self))
41    }
42}
43
44impl Default for WorkerPlugin {
45    fn default() -> Self {
46        WorkerPlugin::None
47    }
48}
49
50#[derive(Serialize, Deserialize)]
51pub struct PluginSettings {
52    pub name: String,
53    pub path: String,
54    pub params: HashMap<String, String>,
55}
56
57impl PluginSettings {
58    pub fn new(
59        name: String,
60        path: String,
61        params: HashMap<String, String>,
62    ) -> Self {
63        Self {
64            name,
65            path,
66            params,
67        }
68    }
69
70    pub fn empty() -> Self {
71        PluginSettings::new("".into(), "".into(), HashMap::new())
72    }
73}
74
75fn plugin_settings(plugin: WorkerPlugin) -> PluginSettings {
76    match plugin {
77        WorkerPlugin::Basic => {
78            PluginSettings::new(
79                WorkerPlugin::as_str(plugin).to_string(),
80                env::full_path(
81                    "$PATOKA_X_DIR/build/src/plugin/basic_plugin.js",
82                    "$PATOKA_X_DIR",
83                    &PATOKA_X_DIR,
84                ),
85                HashMap::new(),
86            )
87        },
88        WorkerPlugin::HeadlessBrowser => {
89            PluginSettings::new(
90                WorkerPlugin::as_str(plugin).to_string(),
91                env::full_path(
92                    "$PATOKA_X_DIR/build/src/plugin/headless_browser_plugin.js",
93                    "$PATOKA_X_DIR",
94                    &PATOKA_X_DIR,
95                ),
96                params_headless_browser(),
97            )
98        },
99        WorkerPlugin::None => {
100            PluginSettings::empty()
101        },
102    }
103}
104
105pub fn setup_plugin_message(
106    plugin: WorkerPlugin,
107    worker_id: &str,
108) -> WorkerMessage {
109    let settings = plugin_settings(plugin);
110    let data = json!({
111        "plugin": serde_json::to_value(settings).unwrap(),
112    });
113
114    let dest = Dest::Worker;
115
116    let payload = WorkerMessagePayload {
117        dest,
118        worker_id: worker_id.to_string(),
119        task_uuid: String::new(),
120        plugin: WorkerPlugin::as_str(plugin).to_string(),
121        data,
122    };
123
124    WorkerMessage::new(payload)
125}
126
127fn params_headless_browser() -> HashMap<String, String> {
128    let mut params = HashMap::new();
129
130    // User-Agent header
131    params.insert("user_agent".to_string(), user_agent::random_ua());
132
133    // Proxy
134    if let Some(proxy) = proxy::next() {
135        let proxy_server = proxy.type_ + "://" + &proxy.address;
136        params.insert("proxy_server".to_string(), proxy_server);
137    }
138
139    // DevTools
140    if let Some(v) = env::get_opt_var("plugin.headless_browser.dev_tools") {
141        if v == "yes" {
142            params.insert("dev_tools".to_string(), "yes".to_string());
143        }
144    }
145
146    params
147}
patoka/worker/plugin.rs

patoka/worker/
plugin.rs
