diff --git a/src/Filter.cpp b/src/Filter.cpp index df77b5c..24e7eee 100644 --- a/src/Filter.cpp +++ b/src/Filter.cpp @@ -8,7 +8,46 @@ #include "Filter.h" Filter::Filter() { - // TODO Auto-generated constructor stub - +} +Filter::Filter(macAddr a) { + switchMac = a; +} +Filter::Filter(Packet::OpCode o) { + oc = o; } +Filter Filter::opcode(Packet::OpCode o) { + oc = o; + return *this; +} +Filter Filter::mac(macAddr a) { + switchMac = a; + return *this; +} +Filter Filter::tokenid(short i) { + tokenId = i; + return *this; +} + +bool Filter::pass(const Packet p) const { + macAddr none { 0, 0, 0, 0, 0, 0 }; + if (this->oc != Packet::NONE && this->oc != p.getOpCode()) + return false; + if (this->switchMac != none && this->switchMac != p.getSwitchMac()) + return false; + if (this->tokenId != -1 && this->tokenId != p.getTokenId()) + return false; + /* if (fragmentOffset != p.getF && fragmentOffset != -1) + return false;*/ + return true; +} + +bool Filter::operator<(const Filter f) const { + if (this->oc != f.oc) + return this->oc < f.oc; + if (this->switchMac != f.switchMac) + return this->switchMac < f.switchMac; + if (this->tokenId < f.tokenId) + return this->tokenId < f.tokenId; + return false; +} diff --git a/src/Filter.h b/src/Filter.h index cbd7ee9..a9d9213 100644 --- a/src/Filter.h +++ b/src/Filter.h @@ -8,9 +8,23 @@ #ifndef FILTER_H_ #define FILTER_H_ +#include "Packet.h" + class Filter { public: Filter(); + Filter(Packet::OpCode); + Filter(macAddr); + Filter opcode(Packet::OpCode); + Filter mac(macAddr); + Filter tokenid(short); + bool pass(const Packet) const; + bool operator<(const Filter) const; +private: + Packet::OpCode oc = Packet::NONE; + macAddr switchMac = { 0, 0, 0, 0, 0, 0 }; + short tokenId = -1; + short fragmentOffset = -1; }; #endif /* FILTER_H_ */ diff --git a/src/Interactive.cpp b/src/Interactive.cpp index 50a3bbc..19ab930 100644 --- a/src/Interactive.cpp +++ b/src/Interactive.cpp @@ -22,14 +22,8 @@ using namespace std; int Interactive::loop() { string cmd; vector v; - //vector vc; - //const char** argv; - //int argc; - Program p = Program(); - p.init(); - while (1) { cmd = readline("smrtlink> "); if (!cmd.compare("quit") || !cmd.compare("q")) @@ -37,27 +31,16 @@ int Interactive::loop() { if (!cmd.empty()) { add_history(cmd.c_str()); v = boost::program_options::split_unix(cmd); - - //vc = vector(v.size()); - //std::transform(begin(v), end(v), begin(vc), [](std::string const &s) { return s.c_str(); }); - //argv = &vc[0]; - //argc = v.size(); p.run(v); - } - } return 0; } int Interactive::single(vector v) { - Program p = Program(); - p.init(); - p.run(v); - return 0; } diff --git a/src/Packet.h b/src/Packet.h index 395c86d..a34e1c9 100644 --- a/src/Packet.h +++ b/src/Packet.h @@ -13,12 +13,13 @@ #include "Types.h" -static short sequenceId=0; +static short sequenceId = 0; -class Packet -{ +class Packet { public: - enum OpCode {DISCOVERY, GET, REPLY, SET, CONFIRM, NONE}; + enum OpCode { + DISCOVERY, GET, REPLY, SET, CONFIRM, NONE + }; Packet(OpCode); static void encode(bytes&); bytes getBytes(); @@ -37,7 +38,7 @@ public: void setHostMac(macAddr); void setSwitchMac(macAddr); void setCheckSum(int); - void setSequenceId(short ); + void setSequenceId(short); void setPayload(datasets payload); short getTokenId() const; void setTokenId(short tokenId = 0); @@ -50,9 +51,9 @@ private: datasets payload; byte version = 1; byte opCode; - macAddr switchMac { {0, 0, 0, 0, 0, 0}}; - macAddr hostMac { {0, 0, 0, 0, 0, 0}}; - macAddr broadcastMac { {0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF}}; + macAddr switchMac { { 0, 0, 0, 0, 0, 0 } }; + macAddr hostMac { { 0, 0, 0, 0, 0, 0 } }; + macAddr broadcastMac { { 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF } }; short tokenId = 0; short fragmentOffset = 0; int errorCode = 0; @@ -60,7 +61,7 @@ private: short flag = 0; void buildHead(); void buildBody(); - void push(bytes&, int&, short ); + void push(bytes&, int&, short); void push(bytes&, int&, int); void push(bytes&, int&, byte); void push(bytes&, int&, bytes); @@ -70,7 +71,7 @@ private: void pull(bytes&, int&, short &); void pull(bytes&, int&, int&); void pull(bytes&, int&, byte&); - void pull(bytes&, int&, bytes&, unsigned ); + void pull(bytes&, int&, bytes&, unsigned); void pull(bytes&, int&, ipAddr&); void pull(bytes&, int&, macAddr&); void pull(bytes&, int&, dataset&); diff --git a/src/Program.cpp b/src/Program.cpp index ae6c279..d4ce237 100644 --- a/src/Program.cpp +++ b/src/Program.cpp @@ -17,6 +17,7 @@ #include "Socket.h" #include "Switch.h" #include "Packet.h" +#include "Filter.h" #include "Types.h" #include "lookup.h" #include "table.h" @@ -220,18 +221,16 @@ int Program::sniff() { Socket s(io_service); s.setHostIp(host.getIp()); s.init(DST_PORT, SRC_PORT); - s.callback = [](Packet p) { + s.listen([](Packet p) { cout << p.opCodeToString() << "\n"; printHeader(p); printPacket(p); return 0; - }; - s.receive(); + }); io_service.run(); } catch (exception& e) { cerr << "Exception: " << e.what() << "\n"; } - return 0; } @@ -406,7 +405,7 @@ int Program::discover(function c) { Packet p = Packet(Packet::DISCOVERY); p.setHostMac(host.getMac()); p.setPayload( { }); - sock->callback = c; + sock->listen(c,Filter(Packet::REPLY)); sock->send(p); return 0; } @@ -416,7 +415,7 @@ int Program::get(Packet l, datasets t, function c) { p.setSwitchMac(l.getSwitchMac()); p.setHostMac(host.getMac()); p.setPayload(t); - sock->callback = c; + sock->listen(c,Filter(Packet::REPLY)); sock->send(p); return 0; } @@ -431,7 +430,7 @@ int Program::set(Packet l, datasets t, function c) { datasets ld = { { LOGIN_USER, (short) (n.size()), n }, { LOGIN_PASSWORD, (short) (w.size()), w } }; p.setPayload(ld + t); - sock->callback = c; + sock->listen(c,Filter(Packet::CONFIRM)); sock->send(p); return 0; } diff --git a/src/Socket.cpp b/src/Socket.cpp index 74b2a27..97b3596 100644 --- a/src/Socket.cpp +++ b/src/Socket.cpp @@ -5,11 +5,14 @@ * Author: jdi */ #include -#include + +#include + #include #include "Socket.h" #include "Packet.h" #include "Constant.h" +#include "Filter.h" #include "Host.h" #include "Types.h" @@ -17,6 +20,10 @@ Socket::Socket(boost::asio::io_service& io_service) : send_socket_(io_service), receive_socket_(io_service), timer(io_service) { } //, resolver( io_service) + +Socket::~Socket() { +} + void Socket::init(short dst_port, short src_port) { if (options.flags.REVERSE) { short p = dst_port; @@ -50,6 +57,15 @@ void Socket::setHostIp(ipAddr ip) { local_ip = ip; } +void Socket::listen(Listener l, Filter f) { + if (callback.find(f) == callback.end()) { + callback.insert(ListenerPair(f, l)); + } else { + callback[f] = l; + } + receive(); +} + void Socket::send(Packet p) { bytes data = p.getBytes(); p.encode(data); @@ -76,11 +92,15 @@ void Socket::receive() { data.resize(bytes_recvd); Packet p = Packet(Packet::NONE); p.encode(data); - // std::cout << "err" << p.getErrorCode() < #include #include "Packet.h" +#include "Filter.h" #include "Types.h" #define SRC_PORT 29809 @@ -18,22 +20,19 @@ #define MAX_LENGTH 1024 typedef std::function Listener; +typedef std::pair ListenerPair; class Socket { public: Socket(boost::asio::io_service&); - virtual ~Socket() { - } + virtual ~Socket(); void init(short, short); void send(Packet); void setHostIp(ipAddr); - void listen(); - Listener callback = [](Packet a) { - return 0; - }; - void receive(); + void listen(Listener l, Filter f = Filter()); private: + void receive(); void settimeout(); boost::asio::ip::udp::socket send_socket_; boost::asio::ip::udp::socket receive_socket_; @@ -44,7 +43,7 @@ private: boost::asio::deadline_timer timer; bytes data = bytes(MAX_LENGTH); ipAddr local_ip; - int initialized = 0; + std::map callback = { }; }; diff --git a/src/Types.h b/src/Types.h index e236c10..1d0a6c5 100644 --- a/src/Types.h +++ b/src/Types.h @@ -50,6 +50,20 @@ public: else break; } } + + bool operator==(const macAddr &A) { + for (unsigned i = 0; i < 6; i++) { + if(A[i]!=(*this)[i])return false; + } + return true; + } + + bool operator!=(const macAddr &A) { + for (unsigned i = 0; i < 6; i++) { + if(A[i]!=(*this)[i])return true; + } + return false; + } }; /*