Ensure consistent column order.
@@ -33,10 +33,12 @@ class HostFootprint():
self.raw_args = raw_args
@staticmethod
- def drop_cols(df):
+ def reorder_drop_cols(df):
# TODO: need host_key and tshark_srcips to send source_ip/source_mac to Poseidon.
cols = [col for col in ('host_key', 'tshark_srcips') if col in df.columns]
- return df.drop(columns=cols)
+ df = df.drop(columns=cols)
+ # Dataframe column order must be the same for train/predict!
+ return df.reindex(columns=sorted(df.columns))
@staticmethod
def serialize_label_encoder(le, path):
@@ -122,7 +124,7 @@ class HostFootprint():
# Load data from host footprint .csv
df = pd.read_csv(self.path)
- df = self.drop_cols(df)
+ df = self.reorder_drop_cols(df)
df = df.fillna(0)
# Split dataframe into X (the input features or predictors)
@@ -190,7 +192,7 @@ class HostFootprint():
# Load data from host footprint .csv
df = pd.read_csv(self.path)
- df = self.drop_cols(df)
+ df = self.reorder_drop_cols(df)
# Split dataframe into X (the input features or predictors)
# and y (the target or outcome or dependent variable)
